gsplat 1.0.0

data/ext/gsplat_native/raster_forward.c

@@ -0,0 +1,150 @@
+#include <math.h>
+#include <string.h>
+#include "common.h"
+#include "rasterization.h"
+
+void gs_init_rasterization_backward(VALUE native);
+
+static size_t
+gs_pixel_index(const gs_raster_forward_args *args, size_t camera, int x, int y)
+{
+    return (camera * args->height + (size_t)y) * args->width + (size_t)x;
+}
+
+static void
+gs_raster_pixel(
+    const gs_raster_forward_args *args,
+    size_t camera,
+    int x,
+    int y,
+    size_t range_start,
+    size_t range_end
+)
+{
+    size_t pixel = gs_pixel_index(args, camera, x, y);
+    float *output = args->render_colors + pixel * args->channels;
+    float transmittance = 1.0f;
+    for (size_t intersection = range_start; intersection < range_end; ++intersection) {
+        int32_t gaussian = args->flatten_ids[intersection];
+        float delta_x = args->means[gaussian * 2] - ((float)x + 0.5f);
+        float delta_y = args->means[gaussian * 2 + 1] - ((float)y + 0.5f);
+        const float *conic = args->conics + gaussian * 3;
+        float sigma = 0.5f * (
+            conic[0] * delta_x * delta_x + conic[2] * delta_y * delta_y
+        ) + conic[1] * delta_x * delta_y;
+        float alpha = args->opacities[gaussian] * expf(-sigma);
+        if (alpha > GS_ALPHA_CLAMP) alpha = GS_ALPHA_CLAMP;
+        if (sigma < 0.0f || alpha < GS_ALPHA_SKIP) continue;
+        float next_transmittance = transmittance * (1.0f - alpha);
+        if (next_transmittance <= GS_TRANSMITTANCE_STOP) break;
+        float visibility = alpha * transmittance;
+        const float *color = args->colors + gaussian * args->channels;
+        for (size_t channel = 0; channel < args->channels; ++channel) {
+            output[channel] += visibility * color[channel];
+        }
+        transmittance = next_transmittance;
+        args->last_ids[pixel] = (int32_t)intersection;
+    }
+    if (args->backgrounds != NULL) {
+        const float *background = args->backgrounds + camera * args->channels;
+        for (size_t channel = 0; channel < args->channels; ++channel) {
+            output[channel] += transmittance * background[channel];
+        }
+    }
+    args->render_alphas[pixel] = 1.0f - transmittance;
+}
+
+void *
+gs_raster_forward_without_gvl(void *opaque)
+{
+    gs_raster_forward_args *args = (gs_raster_forward_args *)opaque;
+    size_t tiles_per_camera = (size_t)args->tile_width * args->tile_height;
+    size_t total_tiles = args->cameras * tiles_per_camera;
+#ifdef GSPLAT_OPENMP
+#pragma omp parallel for schedule(dynamic, 1)
+#endif
+    for (size_t tile = 0; tile < total_tiles; ++tile) {
+        size_t camera = tile / tiles_per_camera;
+        size_t local_tile = tile % tiles_per_camera;
+        int tile_y = (int)(local_tile / args->tile_width);
+        int tile_x = (int)(local_tile % args->tile_width);
+        int x_start = tile_x * args->tile_size;
+        int y_start = tile_y * args->tile_size;
+        int x_end = x_start + args->tile_size;
+        int y_end = y_start + args->tile_size;
+        if (x_end > args->width) x_end = args->width;
+        if (y_end > args->height) y_end = args->height;
+        if (x_start >= args->width || y_start >= args->height) continue;
+        size_t range_start = (size_t)args->offsets[tile];
+        size_t range_end = tile + 1 < total_tiles ?
+            (size_t)args->offsets[tile + 1] : args->intersection_count;
+        int masked = args->masks != NULL && args->masks[tile] == 0;
+        if ((masked || range_start == range_end) && args->backgrounds == NULL) continue;
+        for (int y = y_start; y < y_end; ++y) {
+            for (int x = x_start; x < x_end; ++x) {
+                if (masked) {
+                    size_t pixel = gs_pixel_index(args, camera, x, y);
+                    if (args->backgrounds != NULL) {
+                        memcpy(
+                            args->render_colors + pixel * args->channels,
+                            args->backgrounds + camera * args->channels,
+                            args->channels * sizeof(float)
+                        );
+                    }
+                } else {
+                    gs_raster_pixel(args, camera, x, y, range_start, range_end);
+                }
+            }
+        }
+    }
+    return NULL;
+}
+
+static VALUE
+gs_raster_forward(int argc, VALUE *argv, VALUE self)
+{
+    (void)self;
+    if (argc != 11) rb_raise(rb_eArgError, "expected 11 raster forward arguments");
+    gs_require_sfloat(argv[0], "means2d");
+    gs_require_sfloat(argv[1], "conics");
+    gs_require_sfloat(argv[2], "colors");
+    gs_require_sfloat(argv[3], "opacities");
+    if (!NIL_P(argv[4])) gs_require_sfloat(argv[4], "backgrounds");
+    if (!NIL_P(argv[5])) gs_require_int32(argv[5], "masks");
+    gs_require_int32(argv[9], "isect_offsets");
+    gs_require_int32(argv[10], "flatten_ids");
+    narray_t *means = gs_narray(argv[0]);
+    narray_t *colors = gs_narray(argv[2]);
+    narray_t *offsets = gs_narray(argv[9]);
+    size_t color_shape[4] = {
+        means->shape[0], (size_t)NUM2INT(argv[7]), (size_t)NUM2INT(argv[6]), colors->shape[2]
+    };
+    size_t alpha_shape[4] = {means->shape[0], color_shape[1], color_shape[2], 1};
+    size_t last_shape[3] = {means->shape[0], color_shape[1], color_shape[2]};
+    VALUE render_colors = nary_new(numo_cSFloat, 4, color_shape);
+    VALUE render_alphas = nary_new(numo_cSFloat, 4, alpha_shape);
+    VALUE last_ids = nary_new(numo_cInt32, 3, last_shape);
+    memset(GS_SFLOAT_WRITE(render_colors), 0, gs_narray(render_colors)->size * sizeof(float));
+    memset(GS_SFLOAT_WRITE(render_alphas), 0, gs_narray(render_alphas)->size * sizeof(float));
+    memset(GS_INT32_WRITE(last_ids), 0, gs_narray(last_ids)->size * sizeof(int32_t));
+    gs_raster_forward_args args = {
+        GS_SFLOAT_READ(argv[0]), GS_SFLOAT_READ(argv[1]), GS_SFLOAT_READ(argv[2]),
+        GS_SFLOAT_READ(argv[3]), NIL_P(argv[4]) ? NULL : GS_SFLOAT_READ(argv[4]),
+        NIL_P(argv[5]) ? NULL : GS_INT32_READ(argv[5]), GS_INT32_READ(argv[9]),
+        GS_INT32_READ(argv[10]), GS_SFLOAT_WRITE(render_colors), GS_SFLOAT_WRITE(render_alphas),
+        GS_INT32_WRITE(last_ids), means->shape[0], means->shape[1], colors->shape[2],
+        gs_narray(argv[10])->size, NUM2INT(argv[6]), NUM2INT(argv[7]), NUM2INT(argv[8]),
+        (int)offsets->shape[2], (int)offsets->shape[1]
+    };
+    rb_thread_call_without_gvl(gs_raster_forward_without_gvl, &args, RUBY_UBF_IO, NULL);
+    RB_GC_GUARD(argv[0]);
+    RB_GC_GUARD(render_colors);
+    return rb_ary_new_from_args(3, render_colors, render_alphas, last_ids);
+}
+
+void
+gs_init_rasterization(VALUE native)
+{
+    rb_define_singleton_method(native, "rasterize_forward_sfloat", gs_raster_forward, -1);
+    gs_init_rasterization_backward(native);
+}

data/ext/gsplat_native/rasterization.h

@@ -0,0 +1,51 @@
+#ifndef GSPLAT_RASTERIZATION_H
+#define GSPLAT_RASTERIZATION_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+#define GS_ALPHA_CLAMP 0.999f
+#define GS_ALPHA_SKIP (1.0f / 255.0f)
+#define GS_TRANSMITTANCE_STOP 1.0e-4f
+
+typedef struct {
+    const float *means;
+    const float *conics;
+    const float *colors;
+    const float *opacities;
+    const float *backgrounds;
+    const int32_t *masks;
+    const int32_t *offsets;
+    const int32_t *flatten_ids;
+    float *render_colors;
+    float *render_alphas;
+    int32_t *last_ids;
+    size_t cameras;
+    size_t gaussians;
+    size_t channels;
+    size_t intersection_count;
+    int width;
+    int height;
+    int tile_size;
+    int tile_width;
+    int tile_height;
+} gs_raster_forward_args;
+
+typedef struct {
+    gs_raster_forward_args forward;
+    const float *render_alphas;
+    const int32_t *last_ids;
+    const float *grad_render_colors;
+    const float *grad_render_alphas;
+    float *grad_means;
+    float *grad_conics;
+    float *grad_colors;
+    float *grad_opacities;
+    float *grad_backgrounds;
+    float *grad_means_abs;
+} gs_raster_backward_args;
+
+void *gs_raster_forward_without_gvl(void *opaque);
+void *gs_raster_backward_without_gvl(void *opaque);
+
+#endif

data/ext/gsplat_native/spherical_harmonics.c

@@ -0,0 +1,129 @@
+#include <math.h>
+#include "common.h"
+
+typedef struct {
+    int degree;
+    size_t count;
+    size_t basis_count;
+    size_t channels;
+    const float *directions;
+    const float *coefficients;
+    float *output;
+} gs_sh_args;
+
+static void
+gs_sh_basis(int degree, float x, float y, float z, float *basis)
+{
+    const float c0 = 0.2820947917738781f;
+    const float c1 = 0.48860251190292f;
+    basis[0] = c0;
+    if (degree < 1) return;
+    basis[1] = -c1 * y;
+    basis[2] = c1 * z;
+    basis[3] = -c1 * x;
+    if (degree < 2) return;
+    basis[4] = 1.092548430592079f * x * y;
+    basis[5] = -1.092548430592079f * y * z;
+    basis[6] = 0.3153915652525201f * (3.0f * z * z - 1.0f);
+    basis[7] = -1.092548430592079f * x * z;
+    basis[8] = 0.5462742152960395f * (x * x - y * y);
+    if (degree < 3) return;
+    basis[9] = -0.5900435899266435f * y * (3.0f * x * x - y * y);
+    basis[10] = 2.890611442640554f * x * y * z;
+    basis[11] = -0.4570457994644658f * y * (5.0f * z * z - 1.0f);
+    basis[12] = 0.3731763325901154f * z * (5.0f * z * z - 3.0f);
+    basis[13] = -0.4570457994644658f * x * (5.0f * z * z - 1.0f);
+    basis[14] = 1.445305721320277f * z * (x * x - y * y);
+    basis[15] = -0.5900435899266435f * x * (x * x - 3.0f * y * y);
+    if (degree < 4) return;
+    float x2 = x * x, y2 = y * y, z2 = z * z;
+    basis[16] = 2.5033429417967046f * x * y * (x2 - y2);
+    basis[17] = -1.770130769779931f * y * z * (3.0f * x2 - y2);
+    basis[18] = 0.9461746957575601f * x * y * (7.0f * z2 - 1.0f);
+    basis[19] = -0.6690465435572892f * y * z * (7.0f * z2 - 3.0f);
+    basis[20] = 0.1057855469152043f * (z2 * (35.0f * z2 - 30.0f) + 3.0f);
+    basis[21] = -0.6690465435572892f * x * z * (7.0f * z2 - 3.0f);
+    basis[22] = 0.47308734787878f * (x2 - y2) * (7.0f * z2 - 1.0f);
+    basis[23] = -1.770130769779931f * x * z * (x2 - 3.0f * y2);
+    basis[24] = 0.6258357354491763f * (x2 * x2 - 6.0f * x2 * y2 + y2 * y2);
+}
+
+static void *
+gs_sh_without_gvl(void *opaque)
+{
+    gs_sh_args *args = (gs_sh_args *)opaque;
+#ifdef GSPLAT_OPENMP
+#pragma omp parallel for
+#endif
+    for (size_t index = 0; index < args->count; ++index) {
+        const float *direction = args->directions + index * 3;
+        float norm = sqrtf(
+            direction[0] * direction[0] +
+            direction[1] * direction[1] +
+            direction[2] * direction[2]
+        );
+        if (norm < 1e-12f) norm = 1e-12f;
+        float basis[25] = {0};
+        gs_sh_basis(
+            args->degree,
+            direction[0] / norm,
+            direction[1] / norm,
+            direction[2] / norm,
+            basis
+        );
+        for (size_t channel = 0; channel < args->channels; ++channel) {
+            float color = 0.0f;
+            for (size_t coefficient = 0; coefficient < args->basis_count; ++coefficient) {
+                size_t offset = (index * args->basis_count + coefficient) * args->channels + channel;
+                color += basis[coefficient] * args->coefficients[offset];
+            }
+            args->output[index * args->channels + channel] = color;
+        }
+    }
+    return NULL;
+}
+
+static VALUE
+gs_sh_forward(VALUE self, VALUE degree_value, VALUE directions, VALUE coefficients)
+{
+    (void)self;
+    gs_require_sfloat(directions, "directions");
+    gs_require_sfloat(coefficients, "coefficients");
+    int degree = NUM2INT(degree_value);
+    narray_t *direction_array = gs_narray(directions);
+    narray_t *coefficient_array = gs_narray(coefficients);
+    if (degree < 0 || degree > 4 || direction_array->shape[direction_array->ndim - 1] != 3) {
+        rb_raise(rb_eArgError, "invalid spherical harmonics inputs");
+    }
+    size_t count = direction_array->size / 3;
+    size_t basis_count = coefficient_array->shape[coefficient_array->ndim - 2];
+    size_t channels = coefficient_array->shape[coefficient_array->ndim - 1];
+    if (basis_count > 25 || basis_count < (size_t)((degree + 1) * (degree + 1)) ||
+        coefficient_array->size != count * basis_count * channels) {
+        rb_raise(rb_eArgError, "invalid spherical harmonics coefficient shape");
+    }
+
+    size_t *shape = ALLOCA_N(size_t, direction_array->ndim);
+    for (int axis = 0; axis < direction_array->ndim - 1; ++axis) {
+        shape[axis] = direction_array->shape[axis];
+    }
+    shape[direction_array->ndim - 1] = channels;
+    VALUE output = nary_new(numo_cSFloat, direction_array->ndim, shape);
+    gs_sh_args args = {
+        degree, count, basis_count, channels,
+        GS_SFLOAT_READ(directions),
+        GS_SFLOAT_READ(coefficients),
+        GS_SFLOAT_WRITE(output)
+    };
+    rb_thread_call_without_gvl(gs_sh_without_gvl, &args, RUBY_UBF_IO, NULL);
+    RB_GC_GUARD(directions);
+    RB_GC_GUARD(coefficients);
+    RB_GC_GUARD(output);
+    return output;
+}
+
+void
+gs_init_spherical_harmonics(VALUE native)
+{
+    rb_define_singleton_method(native, "spherical_harmonics_forward_sfloat", gs_sh_forward, 3);
+}

data/gsplat.gemspec

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require_relative "lib/gsplat/version"
+
+Gem::Specification.new do |spec|
+  spec.name = "gsplat"
+  spec.version = Gsplat::VERSION
+  spec.authors = ["Yudai Takada"]
+  spec.email = ["t.yudai92@gmail.com"]
+
+  spec.summary = "Differentiable 3D Gaussian splatting for Ruby"
+  spec.description = "A CPU differentiable 3D Gaussian splatting renderer, trainer, and IO toolkit " \
+                     "with Numo::NArray and optional OpenMP acceleration."
+  spec.homepage = "https://github.com/ydah/gsplat"
+  spec.license = "Apache-2.0"
+  spec.required_ruby_version = ">= 3.2.0"
+  spec.metadata["rubygems_mfa_required"] = "true"
+  spec.metadata["source_code_uri"] = spec.homepage
+  spec.metadata["documentation_uri"] = "https://rubydoc.info/gems/gsplat/#{spec.version}"
+
+  release_files = %w[LICENSE.txt README.md gsplat.gemspec]
+  release_roots = %w[docs/ examples/ ext/ lib/]
+  spec.files = IO.popen(%w[git ls-files -z], chdir: __dir__, err: IO::NULL) do |ls|
+    ls.readlines("\x0", chomp: true).select do |file|
+      release_files.include?(file) || file.start_with?(*release_roots)
+    end
+  end
+  spec.files |= release_files.select { |file| File.file?(File.join(__dir__, file)) }
+  spec.extra_rdoc_files = %w[README.md LICENSE.txt]
+  spec.bindir = "exe"
+  spec.executables = spec.files.grep(%r{\Aexe/}) { |f| File.basename(f) }
+  spec.extensions = ["ext/gsplat_native/extconf.rb"]
+  spec.require_paths = ["lib"]
+
+  spec.add_dependency "numo-narray", ">= 0.9", "< 1.0"
+end

data/lib/gsplat/autograd/context.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+module Gsplat
+  # Lightweight reverse-mode differentiation for coarse gsplat operations.
+  module Autograd
+    # Thread-local key used to suspend graph recording.
+    # @api private
+    GRAD_ENABLED_KEY = :gsplat_autograd_grad_enabled
+
+    # Per-operation values retained for the backward pass.
+    class Context
+      attr_reader :inputs, :needs_input_grad, :saved_values
+
+      # @param needs_input_grad [Array<Boolean>] gradient requirement for each input
+      # @param inputs [Array<Object, Variable>] original Function inputs
+      def initialize(needs_input_grad, inputs = [])
+        @needs_input_grad = needs_input_grad.freeze
+        @inputs = inputs
+        @saved_values = []
+      end
+
+      # Saves values required by the operation's backward implementation.
+      #
+      # @param values [Array<Object>]
+      # @return [void]
+      def save(*values)
+        @saved_values.concat(values)
+      end
+
+      # Releases references retained for backward.
+      #
+      # @return [void]
+      def clear
+        @saved_values.clear
+        @inputs = []
+      end
+    end
+
+    module_function
+
+    # Runs a block without recording an autograd graph.
+    #
+    # @yieldreturn [Object]
+    def no_grad
+      thread = Thread.current
+      previous = thread.thread_variable_get(GRAD_ENABLED_KEY)
+      thread.thread_variable_set(GRAD_ENABLED_KEY, false)
+      yield
+    ensure
+      thread.thread_variable_set(GRAD_ENABLED_KEY, previous)
+    end
+
+    # Whether new operations should attach backward nodes.
+    #
+    # @return [Boolean]
+    def grad_enabled?
+      Thread.current.thread_variable_get(GRAD_ENABLED_KEY) != false
+    end
+  end
+end

data/lib/gsplat/autograd/function.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Autograd
+    # A recorded invocation shared by all outputs of one Function.
+    class GraphNode
+      attr_reader :function, :context, :inputs
+      attr_accessor :outputs
+
+      def initialize(function, context, inputs)
+        @function = function
+        @context = context
+        @inputs = inputs
+        @outputs = []
+      end
+
+      # Delegates output gradients to the operation's backward method.
+      # @api private
+      def backward(*grad_outputs)
+        function.backward(context, *grad_outputs)
+      end
+
+      # Releases graph references after a backward traversal.
+      # @api private
+      def release
+        outputs.each { |output| output.clear_creator(self) }
+        context.clear
+        @inputs = []
+        @outputs = []
+      end
+    end
+
+    # Base class for differentiable coarse-grained operations.
+    class Function
+      class << self
+        # Executes forward and records a graph node when any input needs a gradient.
+        #
+        # @param inputs [Array<Object, Variable>]
+        # @return [Variable, Array<Variable>]
+        def apply(*inputs, **)
+          needs_input_grad = inputs.map { |input| input.is_a?(Variable) && input.requires_grad? }
+          context = Context.new(needs_input_grad, inputs)
+          raw_inputs = inputs.map { |input| input.is_a?(Variable) ? input.data : input }
+          result = forward(context, *raw_inputs, **)
+          multiple_outputs = result.is_a?(Array)
+          output_data = multiple_outputs ? result : [result]
+          requires_grad = Autograd.grad_enabled? && needs_input_grad.any?
+          node = GraphNode.new(self, context, inputs) if requires_grad
+          outputs = output_data.map do |data|
+            validate_output!(data)
+            Variable.new(data, requires_grad: requires_grad, creator: node)
+          end
+          node.outputs = outputs if node
+
+          multiple_outputs ? outputs : outputs.first
+        end
+
+        private
+
+        def validate_output!(data)
+          return if data.is_a?(Numo::NArray)
+
+          raise Gsplat::Error, "Function.forward must return Numo::NArray output(s), got #{data.class}"
+        end
+      end
+    end
+  end
+end

data/lib/gsplat/autograd/variable.rb

@@ -0,0 +1,159 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Autograd
+    # A Numo array with an optional reverse-mode gradient and creator node.
+    class Variable
+      attr_reader :absgrad, :data, :grad, :creator
+
+      # @param data [Numo::NArray] tensor data
+      # @param requires_grad [Boolean] whether to accumulate a gradient
+      # @param creator [GraphNode, nil] producing operation
+      def initialize(data, requires_grad: false, creator: nil)
+        raise ArgumentError, "Variable data must be a Numo::NArray" unless data.is_a?(Numo::NArray)
+
+        @data = data
+        @requires_grad = requires_grad
+        @creator = creator
+        @grad = nil
+        @absgrad = nil
+        @absgrad_targets = []
+      end
+
+      # Whether this variable accumulates gradients.
+      #
+      # @return [Boolean]
+      def requires_grad?
+        @requires_grad
+      end
+
+      # Clears an accumulated leaf gradient.
+      #
+      # @return [void]
+      def zero_grad!
+        @grad = nil
+        @absgrad = nil
+        @absgrad_targets.each { |target, key| target[key] = nil }
+      end
+
+      # Accumulates an auxiliary absolute gradient produced by rasterization.
+      #
+      # @api private
+      # @param gradient [Numo::NArray]
+      # @return [void]
+      def accumulate_absgrad(gradient)
+        gradient = cast_gradient(gradient)
+        unless gradient.shape == data.shape
+          raise ShapeError, "absgrad shape mismatch: expected #{data.shape.inspect}, got #{gradient.shape.inspect}"
+        end
+
+        @absgrad = @absgrad ? @absgrad + gradient : gradient.dup
+        @absgrad_targets.each { |target, key| target[key] = @absgrad }
+      end
+
+      # Binds auxiliary absolute-gradient updates to a metadata hash entry.
+      #
+      # @api private
+      # @return [void]
+      def bind_absgrad(target, key)
+        @absgrad_targets << [target, key]
+        target[key] = @absgrad
+      end
+
+      # Replaces parameter storage after a first-axis structural edit.
+      #
+      # @api private
+      # @param value [Numo::NArray]
+      # @return [void]
+      def replace_data!(value)
+        raise ArgumentError, "replacement data must be a Numo::NArray" unless value.is_a?(Numo::NArray)
+
+        @data = value
+        zero_grad!
+      end
+
+      # Removes a creator after its graph node has completed backward.
+      #
+      # @api private
+      # @param node [GraphNode]
+      # @return [void]
+      def clear_creator(node)
+        @creator = nil if @creator.equal?(node)
+      end
+
+      # Runs reverse-mode differentiation from this variable.
+      #
+      # @param gradient [Numo::NArray, Numeric, nil] output gradient; nil is valid for scalar output
+      # @return [void]
+      def backward(gradient = nil)
+        raise Gsplat::Error, "cannot call backward on a variable that does not require gradients" unless requires_grad?
+
+        accumulate_grad(initial_gradient(gradient))
+        topological_nodes.reverse_each do |node|
+          grad_outputs = node.outputs.map { |output| output.grad || zeros_like(output.data) }
+          grad_inputs = normalize_grad_inputs(node.backward(*grad_outputs), node.inputs.length)
+          node.inputs.each_with_index do |input, index|
+            next unless input.is_a?(Variable) && node.context.needs_input_grad.fetch(index)
+
+            input.accumulate_grad(grad_inputs.fetch(index))
+          end
+          node.release
+        end
+      end
+
+      protected
+
+      def accumulate_grad(gradient)
+        return if gradient.nil?
+
+        gradient = cast_gradient(gradient)
+        unless gradient.shape == data.shape
+          raise ShapeError, "gradient shape mismatch: expected #{data.shape.inspect}, got #{gradient.shape.inspect}"
+        end
+
+        @grad = @grad ? @grad + gradient : gradient.dup
+      end
+
+      private
+
+      def initial_gradient(gradient)
+        return cast_gradient(gradient) unless gradient.nil?
+        raise Gsplat::Error, "backward on a non-scalar output requires an explicit gradient" unless data.size == 1
+
+        data.shape.empty? ? data.class.cast(1) : data.class.ones(*data.shape)
+      end
+
+      def cast_gradient(gradient)
+        gradient = gradient.data if gradient.is_a?(Variable)
+        data.class.cast(gradient)
+      end
+
+      def zeros_like(array)
+        array.shape.empty? ? array.class.cast(0) : array.class.zeros(*array.shape)
+      end
+
+      def normalize_grad_inputs(result, input_count)
+        return result if input_count == 1 && result.is_a?(Array) && result.length == 1
+        return [result] if input_count == 1
+        return result if result.is_a?(Array) && result.length == input_count
+
+        raise Gsplat::Error, "Function.backward returned gradients for the wrong number of inputs"
+      end
+
+      def topological_nodes
+        ordered = []
+        visited = {}.compare_by_identity
+        visit = lambda do |node|
+          return unless node
+          return if visited.key?(node)
+
+          visited[node] = true
+          node.inputs.each { |input| visit.call(input.creator) if input.is_a?(Variable) }
+          ordered << node
+        end
+        visit.call(creator)
+        ordered
+      end
+    end
+  end
+end