gsplat 1.0.0

data/lib/gsplat/optim/selective_adam.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Optim
+    # Adam variant that updates only visible first-axis parameter rows.
+    class SelectiveAdam < Adam
+      # Updates visible rows while leaving hidden parameters and moments intact.
+      #
+      # @param visibility [Numo::Bit, Numo::NArray, Array<Boolean>] mask shaped [N]
+      # @return [SelectiveAdam] self
+      def step(visibility)
+        mask = validate_visibility(visibility)
+        groups.each_value do |group|
+          gradient = group.variable.grad
+          next unless gradient
+
+          update_group(group, gradient, mask)
+        end
+        self
+      end
+
+      private
+
+      # rubocop:disable Metrics/AbcSize
+      def update_group(group, gradient, mask)
+        state = state_for(group.name)
+        next_step = state.step + 1
+        candidate_mean = (beta1 * state.exp_avg) + ((1 - beta1) * gradient)
+        candidate_variance = (beta2 * state.exp_avg_sq) + ((1 - beta2) * (gradient**2))
+        selection = first_axis_selection(mask, gradient.ndim)
+        exp_avg = state.exp_avg.dup
+        exp_avg_sq = state.exp_avg_sq.dup
+        exp_avg[*selection] = candidate_mean[*selection]
+        exp_avg_sq[*selection] = candidate_variance[*selection]
+        corrected_mean = exp_avg / (1 - (beta1**next_step))
+        corrected_variance = exp_avg_sq / (1 - (beta2**next_step))
+        update = group.lr * corrected_mean / ((corrected_variance**0.5) + group.eps)
+        group.variable.data[*selection] = group.variable.data[*selection] - update[*selection]
+        @states[group.name] = State.new(
+          step: next_step,
+          exp_avg: exp_avg,
+          exp_avg_sq: exp_avg_sq
+        )
+      end
+      # rubocop:enable Metrics/AbcSize
+
+      def validate_visibility(visibility)
+        mask = Numo::Bit.cast(visibility)
+        expected = groups.values.first.variable.data.shape[0]
+        unless mask.ndim == 1 && mask.shape[0] == expected
+          raise ShapeError, "expected visibility [#{expected}], got #{mask.shape.inspect}"
+        end
+
+        groups.each_value do |group|
+          next if group.variable.data.shape[0] == expected
+
+          raise ShapeError, "all SelectiveAdam groups must share first-axis size #{expected}"
+        end
+
+        mask
+      end
+
+      def first_axis_selection(mask, dimensions)
+        [mask] + Array.new(dimensions - 1, true)
+      end
+    end
+  end
+end

data/lib/gsplat/rasterization.rb

@@ -0,0 +1,130 @@
+# frozen_string_literal: true
+
+# High-level differentiable rendering API.
+module Gsplat
+  # Rendering pipeline composed from low-level operations.
+  module Rasterization
+    module_function
+
+    # rubocop:disable Metrics/AbcSize, Metrics/MethodLength, Metrics/ParameterLists, Naming/MethodParameterName
+    def render(means:, opacities:, colors:, viewmats:, ks:, width:, height:, quats: nil, scales: nil,
+               covars: nil, near_plane: 0.01, far_plane: 1e10, radius_clip: 0.0, eps2d: 0.3,
+               sh_degree: nil, tile_size: 16, backgrounds: nil, render_mode: "RGB",
+               rasterize_mode: "classic", camera_model: "pinhole", absgrad: false,
+               channel_chunk: 32, packed: false, global_z_order: true,
+               radial_coeffs: nil, tangential_coeffs: nil, thin_prism_coeffs: nil,
+               with_eval3d: false, return_normals: false)
+      # rubocop:enable Metrics/ParameterLists, Naming/MethodParameterName
+      RasterizationValidation.validate_options!(
+        render_mode, rasterize_mode, sh_degree, tile_size, channel_chunk, width, height
+      )
+      RasterizationValidation.warn_unsupported_options(packed, global_z_order)
+      RasterizationValidation.validate_scene!(
+        means, quats, scales, covars, opacities, colors, viewmats, ks, sh_degree
+      )
+      RasterizationValidation.validate_eval3d!(
+        with_eval3d, return_normals, covars, rasterize_mode
+      )
+      camera_count = Ops::TensorOps.data(viewmats).shape[0]
+      calculate_compensations = rasterize_mode == "antialiased"
+      radii, means2d, depths, conics, compensations = Gsplat.fully_fused_projection(
+        means,
+        viewmats: viewmats,
+        ks: ks,
+        width: width,
+        height: height,
+        covars: covars,
+        quats: quats,
+        scales: scales,
+        eps2d: eps2d,
+        near_plane: near_plane,
+        far_plane: far_plane,
+        radius_clip: radius_clip,
+        calc_compensations: calculate_compensations,
+        camera_model: camera_model,
+        radial_coeffs: radial_coeffs,
+        tangential_coeffs: tangential_coeffs,
+        thin_prism_coeffs: thin_prism_coeffs,
+        global_z_order: global_z_order
+      )
+      projected_colors = RasterizationHelpers.prepare_colors(
+        means, colors, viewmats, radii, camera_count, sh_degree
+      )
+      projected_opacities = RasterizationHelpers.camera_broadcast(opacities, camera_count)
+      if calculate_compensations
+        projected_opacities = Ops::TensorOps.apply(
+          Ops::Multiply,
+          projected_opacities,
+          compensations
+        )
+      end
+      raster_features, depth_index = RasterizationHelpers.render_features(
+        projected_colors, depths, render_mode
+      )
+      tile_width = (width.to_f / tile_size).ceil
+      tile_height = (height.to_f / tile_size).ceil
+      tiles_per_gauss, isect_ids, flatten_ids = Gsplat.isect_tiles(
+        Ops::TensorOps.data(means2d),
+        Ops::TensorOps.data(radii),
+        Ops::TensorOps.data(depths),
+        tile_size,
+        tile_width,
+        tile_height
+      )
+      isect_offsets = Gsplat.isect_offset_encode(
+        isect_ids, camera_count, tile_width, tile_height
+      )
+      hit_distance = %w[d Ed RGB-d RGB-Ed].include?(render_mode)
+      eval3d_normals = nil
+      render_colors, render_alphas = if with_eval3d || hit_distance
+                                       eval3d_outputs = RasterizationHelpers.rasterize_eval3d_features(
+                                         means, quats, scales, raster_features, projected_opacities,
+                                         viewmats, ks, width, height, tile_size, isect_offsets,
+                                         flatten_ids, backgrounds: backgrounds,
+                                                      camera_model: camera_model,
+                                                      use_hit_distance: hit_distance,
+                                                      return_normals: return_normals
+                                       )
+                                       eval3d_normals = eval3d_outputs[2] if return_normals
+                                       eval3d_outputs.first(2)
+                                     else
+                                       RasterizationHelpers.rasterize_features(
+                                         means2d, conics, raster_features, projected_opacities,
+                                         width, height, tile_size, isect_offsets, flatten_ids,
+                                         backgrounds: backgrounds, channel_chunk: channel_chunk,
+                                         absgrad: absgrad
+                                       )
+                                     end
+      if %w[Ed ED RGB-Ed RGB+ED].include?(render_mode)
+        render_colors = Ops::TensorOps.apply(
+          Ops::NormalizeDepth,
+          render_colors,
+          render_alphas,
+          depth_index: depth_index
+        )
+      end
+      meta = RasterizationHelpers.metadata(
+        radii, means2d, depths, conics, projected_opacities, tile_width, tile_height,
+        tiles_per_gauss, isect_ids, flatten_ids, isect_offsets, width, height, tile_size,
+        camera_count
+      )
+      meta[:normals] = eval3d_normals if return_normals
+      means2d.bind_absgrad(meta, :means2d_absgrad) if absgrad && means2d.is_a?(Autograd::Variable)
+      [render_colors, render_alphas, meta]
+    end
+    # rubocop:enable Metrics/AbcSize, Metrics/MethodLength
+  end
+
+  class << self
+    # Renders a dense batch of cameras with differentiable 3D Gaussians.
+    #
+    # Geometry uses float32/float64 arrays `means [N,3]`, `quats [N,4]`,
+    # `scales [N,3]`; colors are `[N,D]` or SH `[N,K,D]`, views are
+    # `[C,4,4]`, and intrinsics are `[C,3,3]`. Activated opacities are `[N]`.
+    #
+    # @return [Array] rendered colors `[C,H,W,X]`, alphas `[C,H,W,1]`, and metadata Hash
+    def rasterization(**)
+      Rasterization.render(**)
+    end
+  end
+end

data/lib/gsplat/rasterization_2dgs.rb

@@ -0,0 +1,140 @@
+# frozen_string_literal: true
+
+# Public 2D Gaussian splatting API.
+module Gsplat
+  # High-level 2D Gaussian rendering composition.
+  module Rasterization2DGS
+    module_function
+
+    # rubocop:disable Metrics/ParameterLists, Naming/MethodParameterName
+    def render(means:, quats:, scales:, opacities:, colors:, viewmats:, ks:, width:, height:,
+               near_plane: 0.01, far_plane: 1e10, radius_clip: 0.0, eps2d: 0.3,
+               sh_degree: nil, packed: false, tile_size: 16, backgrounds: nil,
+               render_mode: "RGB", sparse_grad: false, absgrad: false, distloss: false,
+               depth_mode: "expected", rasterize_mode: "classic")
+      # rubocop:enable Metrics/ParameterLists, Naming/MethodParameterName
+      validate_options!(sparse_grad, depth_mode, rasterize_mode)
+      primary, alphas, meta = Gsplat.rasterization(
+        means: means, quats: quats, scales: scales, opacities: opacities, colors: colors,
+        viewmats: viewmats, ks: ks, width: width, height: height, near_plane: near_plane,
+        far_plane: far_plane, radius_clip: radius_clip, eps2d: eps2d, sh_degree: sh_degree,
+        packed: packed, tile_size: tile_size, backgrounds: backgrounds,
+        render_mode: render_mode, absgrad: absgrad, rasterize_mode: rasterize_mode
+      )
+      gaussian_normals = camera_normals(
+        Ops::TensorOps.data(means), Ops::TensorOps.data(quats), Ops::TensorOps.data(viewmats)
+      )
+      render_normals, = Gsplat.rasterization(
+        means: means, quats: quats, scales: scales, opacities: opacities,
+        colors: gaussian_normals, viewmats: viewmats, ks: ks, width: width, height: height,
+        near_plane: near_plane, far_plane: far_plane, radius_clip: radius_clip, eps2d: eps2d,
+        tile_size: tile_size, render_mode: "RGB", absgrad: false
+      )
+      render_median, = Gsplat.rasterization(
+        means: means, quats: quats, scales: scales, opacities: opacities,
+        colors: gaussian_normals, viewmats: viewmats, ks: ks, width: width, height: height,
+        near_plane: near_plane, far_plane: far_plane, radius_clip: radius_clip, eps2d: eps2d,
+        tile_size: tile_size, render_mode: "ED", absgrad: false
+      )
+      surface_normals = normalized_normals(Ops::TensorOps.data(render_normals), Ops::TensorOps.data(alphas))
+      render_distort = distortion_image(
+        Ops::TensorOps.data(render_median), Ops::TensorOps.data(alphas), distloss
+      )
+      meta = two_d_metadata(meta, gaussian_normals, render_distort)
+      [primary, alphas, render_normals, surface_normals, render_distort, render_median, meta]
+    end
+
+    def validate_options!(sparse_grad, depth_mode, rasterize_mode)
+      raise ArgumentError, "sparse_grad requires packed mode and is unsupported" if sparse_grad
+      raise ArgumentError, "2DGS only supports classic rasterization" unless rasterize_mode == "classic"
+      return if %w[expected median].include?(depth_mode.to_s)
+
+      raise ArgumentError, "depth_mode must be \"expected\" or \"median\""
+    end
+    private_class_method :validate_options!
+
+    def camera_normals(means, quats, viewmats)
+      rotations = Math::Quaternion.to_rotmat(quats)
+      world_normals = rotations[true, true, 2]
+      output = means.class.zeros(viewmats.shape[0], means.shape[0], 3)
+      viewmats.shape[0].times do |camera|
+        view_rotation = viewmats[camera, 0...3, 0...3]
+        camera_means = means.dot(view_rotation.transpose) + viewmats[camera, 0...3, 3]
+        normals = world_normals.dot(view_rotation.transpose)
+        facing_away = (normals * camera_means).sum(axis: 1).gt(0)
+        normals[facing_away, true] = -normals[facing_away, true] if facing_away.any?
+        output[camera, true, true] = normals
+      end
+      output
+    end
+    private_class_method :camera_normals
+
+    def normalized_normals(normals, alphas)
+      alpha = alphas[true, true, true, 0]
+      output = normals.class.zeros(*normals.shape)
+      valid = alpha.gt(1e-10)
+      return output unless valid.any?
+
+      expected = normals.dup
+      3.times { |axis| expected[true, true, true, axis][valid] /= alpha[valid] }
+      norm = (expected**2).sum(axis: 3)**0.5
+      valid &= norm.gt(1e-10)
+      3.times { |axis| output[true, true, true, axis][valid] = expected[true, true, true, axis][valid] / norm[valid] }
+      output
+    end
+    private_class_method :normalized_normals
+
+    def distortion_image(median, alphas, enabled)
+      return median.class.zeros(*median.shape) unless enabled
+
+      alpha = alphas[true, true, true, 0]
+      depth = median[true, true, true, 0]
+      output = median.class.zeros(*median.shape)
+      output[true, true, true, 0] = (depth * (1 - alpha)).abs
+      output
+    end
+    private_class_method :distortion_image
+
+    def two_d_metadata(meta, normals, distortion)
+      ray_transforms = normals.class.zeros(*(normals.shape[0...-1] + [3, 3]))
+      ray_transforms[true, true, true, true] = normals.class.eye(3)
+      meta.merge(
+        ray_transforms: ray_transforms,
+        normals: normals,
+        render_distort: distortion,
+        gradient_2dgs: meta.fetch(:means2d)
+      )
+    end
+    private_class_method :two_d_metadata
+  end
+
+  class << self
+    # Rasterizes oriented 2D Gaussian surfels.
+    #
+    # @return [Array] colors, alphas, normals, surface normals, distortion, median depth, metadata
+    def rasterization_2dgs(**)
+      Rasterization2DGS.render(**)
+    end
+
+    # Projection metadata used by the 2DGS rasterizer.
+    # rubocop:disable Metrics/ParameterLists, Naming/MethodParameterName
+    def fully_fused_projection_2dgs(means, quats:, scales:, viewmats:, ks:, width:, height:,
+                                    eps2d: 0.3, near_plane: 0.01, far_plane: 1e10,
+                                    radius_clip: 0.0)
+      # rubocop:enable Metrics/ParameterLists, Naming/MethodParameterName
+      radii, means2d, depths, conics, = fully_fused_projection(
+        means, quats: quats, scales: scales, viewmats: viewmats, ks: ks, width: width,
+               height: height, eps2d: eps2d, near_plane: near_plane, far_plane: far_plane,
+               radius_clip: radius_clip
+      )
+      normals = Rasterization2DGS.send(
+        :camera_normals, Ops::TensorOps.data(means), Ops::TensorOps.data(quats),
+        Ops::TensorOps.data(viewmats)
+      )
+      ray_transforms = Ops::TensorOps.data(conics).class.zeros(
+        *(Ops::TensorOps.data(conics).shape[0...-1] + [3, 3])
+      )
+      [radii, means2d, depths, ray_transforms, normals]
+    end
+  end
+end

data/lib/gsplat/rasterization_helpers.rb

@@ -0,0 +1,162 @@
+# frozen_string_literal: true
+
+module Gsplat
+  # Feature preparation and metadata helpers for the high-level renderer.
+  module RasterizationHelpers
+    module_function
+
+    def camera_broadcast(value, camera_count)
+      shape = Ops::TensorOps.data(value).shape
+      return value if shape[0] == camera_count && shape.length >= 2
+
+      Ops::TensorOps.apply(Ops::CameraBroadcast, value, camera_count)
+    end
+
+    # rubocop:disable Metrics/ParameterLists
+    def prepare_colors(means, colors, viewmats, radii, camera_count, sh_degree)
+      # rubocop:enable Metrics/ParameterLists
+      unless sh_degree
+        return colors if Ops::TensorOps.data(colors).ndim == 3
+
+        return Ops::TensorOps.apply(Ops::CameraBroadcast, colors, camera_count)
+      end
+
+      directions = Ops::TensorOps.apply(Ops::CameraDirections, means, viewmats)
+      coefficients = Ops::TensorOps.apply(Ops::CameraBroadcast, colors, camera_count)
+      evaluated = Gsplat.spherical_harmonics(
+        sh_degree,
+        directions,
+        coefficients,
+        masks: visible_radii(radii)
+      )
+      Ops::TensorOps.apply(Ops::AddClampMin, evaluated, offset: 0.5, minimum: 0.0)
+    end
+
+    def visible_radii(radii)
+      values = Ops::TensorOps.data(radii)
+      return values.gt(0) unless values.ndim >= 3 && values.shape[-1] == 2
+
+      values.min(axis: values.ndim - 1).gt(0)
+    end
+
+    def render_features(colors, depths, render_mode)
+      return [colors, nil] if render_mode == "RGB"
+
+      depth_features = Ops::TensorOps.apply(Ops::DepthFeatures, depths)
+      return [depth_features, 0] if %w[d Ed D ED].include?(render_mode)
+
+      feature_count = Ops::TensorOps.data(colors).shape[-1]
+      [Ops::TensorOps.apply(Ops::ConcatDepth, colors, depth_features), feature_count]
+    end
+
+    # rubocop:disable Metrics/ParameterLists
+    def rasterize_features(means2d, conics, features, opacities, width, height, tile_size, offsets, flatten_ids,
+                           backgrounds:, channel_chunk:, absgrad:)
+      # rubocop:enable Metrics/ParameterLists
+      channel_count = Ops::TensorOps.data(features).shape[-1]
+      if channel_count <= channel_chunk
+        return Gsplat.rasterize_to_pixels(
+          means2d,
+          conics,
+          features,
+          opacities,
+          width,
+          height,
+          tile_size,
+          offsets,
+          flatten_ids,
+          backgrounds: backgrounds,
+          absgrad: absgrad
+        )
+      end
+
+      rasterize_feature_chunks(
+        means2d, conics, features, opacities, width, height, tile_size, offsets, flatten_ids,
+        backgrounds: backgrounds, channel_chunk: channel_chunk, absgrad: absgrad
+      )
+    end
+
+    # rubocop:disable Metrics/ParameterLists
+    def rasterize_feature_chunks(means2d, conics, features, opacities, width, height,
+                                 tile_size, offsets, flatten_ids,
+                                 backgrounds:, channel_chunk:, absgrad:)
+      # rubocop:enable Metrics/ParameterLists
+      channel_count = Ops::TensorOps.data(features).shape[-1]
+      rendered_chunks = []
+      alpha = nil
+      (0...channel_count).step(channel_chunk) do |start|
+        range = start...[start + channel_chunk, channel_count].min
+        feature_chunk = Ops::TensorOps.apply(Ops::FeatureSlice, features, range)
+        background_chunk = backgrounds && Ops::TensorOps.apply(Ops::FeatureSlice, backgrounds, range)
+        rendered, chunk_alpha = Gsplat.rasterize_to_pixels(
+          means2d,
+          conics,
+          feature_chunk,
+          opacities,
+          width,
+          height,
+          tile_size,
+          offsets,
+          flatten_ids,
+          backgrounds: background_chunk,
+          absgrad: absgrad
+        )
+        rendered_chunks << rendered
+        alpha ||= chunk_alpha
+      end
+      [Ops::TensorOps.apply(Ops::ConcatFeatures, *rendered_chunks), alpha]
+    end
+    private_class_method :rasterize_feature_chunks
+
+    # rubocop:disable Metrics/ParameterLists
+    def rasterize_eval3d_features(means, quats, scales, features, opacities, viewmats, intrinsics,
+                                  width, height, tile_size, offsets, flatten_ids, backgrounds:,
+                                  camera_model:, use_hit_distance:, return_normals:)
+      # rubocop:enable Metrics/ParameterLists
+      raise ArgumentError, "world-space rendering requires quats and scales" unless quats && scales
+
+      inputs = [means, quats, scales, features, opacities, backgrounds]
+      options = {
+        viewmats: Ops::TensorOps.data(viewmats),
+        intrinsics: Ops::TensorOps.data(intrinsics),
+        width: width,
+        height: height,
+        tile_size: tile_size,
+        offsets: offsets,
+        flatten_ids: flatten_ids,
+        camera_model: camera_model.to_s,
+        use_hit_distance: use_hit_distance,
+        return_normals: return_normals
+      }
+      return Ops::Eval3dRasterize.apply(*inputs, **options) if inputs.any?(Autograd::Variable)
+
+      Backend::RubyEval3dRasterizer.forward(*inputs, **options)
+    end
+
+    # rubocop:disable Metrics/ParameterLists
+    def metadata(radii, means2d, depths, conics, opacities, tile_width, tile_height,
+                 tiles_per_gauss, isect_ids, flatten_ids, isect_offsets, width, height,
+                 tile_size, camera_count)
+      # rubocop:enable Metrics/ParameterLists
+      {
+        camera_ids: nil,
+        gaussian_ids: nil,
+        radii: radii,
+        means2d: means2d,
+        depths: depths,
+        conics: conics,
+        opacities: opacities,
+        tile_width: tile_width,
+        tile_height: tile_height,
+        tiles_per_gauss: tiles_per_gauss,
+        isect_ids: isect_ids,
+        flatten_ids: flatten_ids,
+        isect_offsets: isect_offsets,
+        width: width,
+        height: height,
+        tile_size: tile_size,
+        n_cameras: camera_count
+      }
+    end
+  end
+end

data/lib/gsplat/rasterization_validation.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+module Gsplat
+  # Input and option validation for the high-level renderer.
+  module RasterizationValidation
+    RENDER_MODES = %w[RGB d Ed D ED RGB-d RGB-Ed RGB+D RGB+ED].freeze
+    RASTERIZE_MODES = %w[classic antialiased].freeze
+
+    module_function
+
+    # rubocop:disable Metrics/ParameterLists
+    def validate_options!(render_mode, rasterize_mode, sh_degree, tile_size, channel_chunk, width, height)
+      # rubocop:enable Metrics/ParameterLists
+      raise ArgumentError, "unsupported render_mode #{render_mode.inspect}" unless RENDER_MODES.include?(render_mode)
+      unless RASTERIZE_MODES.include?(rasterize_mode)
+        raise ArgumentError, "unsupported rasterize_mode #{rasterize_mode.inspect}"
+      end
+
+      unless sh_degree.nil? || (0..4).cover?(sh_degree)
+        raise ArgumentError, "sh_degree must be nil or an integer from 0 through 4"
+      end
+
+      values = { width: width, height: height, tile_size: tile_size, channel_chunk: channel_chunk }
+      invalid = values.find { |_name, value| !value.is_a?(Integer) || !value.positive? }
+      raise ArgumentError, "#{invalid[0]} must be a positive integer" if invalid
+    end
+
+    def warn_unsupported_options(packed, _global_z_order)
+      Gsplat.logger.warn("packed mode is unsupported; using dense mode") if packed
+    end
+
+    def validate_eval3d!(with_eval3d, return_normals, covars, rasterize_mode)
+      raise ArgumentError, "return_normals requires with_eval3d: true" if return_normals && !with_eval3d
+      raise ArgumentError, "with_eval3d requires quats and scales instead of covars" if with_eval3d && covars
+      return unless with_eval3d && rasterize_mode != "classic"
+
+      raise ArgumentError, "with_eval3d requires rasterize_mode: \"classic\""
+    end
+
+    # rubocop:disable Metrics/ParameterLists
+    def validate_scene!(means, quats, scales, covars, opacities, colors, viewmats, intrinsics, sh_degree)
+      # rubocop:enable Metrics/ParameterLists
+      means_shape = data(means).shape
+      unless means_shape.length == 2 && means_shape[-1] == 3
+        raise ShapeError, "expected means [N,3], got #{means_shape.inspect}"
+      end
+
+      gaussian_count = means_shape[0]
+      opacity_shape = data(opacities).shape
+      unless [[gaussian_count], [data(viewmats).shape[0], gaussian_count]].include?(opacity_shape)
+        raise ShapeError, "expected opacities [N] or [C,N], got #{opacity_shape.inspect}"
+      end
+
+      validate_geometry!(quats, scales, covars, gaussian_count)
+      camera_count = data(viewmats).shape[0]
+      unless data(viewmats).shape == [camera_count, 4, 4] && data(intrinsics).shape == [camera_count, 3, 3]
+        raise ShapeError,
+              "expected viewmats [C,4,4] and ks [C,3,3], " \
+              "got #{data(viewmats).shape.inspect} and #{data(intrinsics).shape.inspect}"
+      end
+
+      validate_color_shape!(data(colors).shape, gaussian_count, camera_count, sh_degree)
+    end
+
+    def validate_geometry!(quats, scales, covars, gaussian_count)
+      if covars
+        covar_shape = data(covars).shape
+        return if [[gaussian_count, 3, 3], [gaussian_count, 6]].include?(covar_shape)
+
+        raise ShapeError, "expected covars [N,3,3] or [N,6], got #{covar_shape.inspect}"
+      end
+      return if data(quats).shape == [gaussian_count, 4] && data(scales).shape == [gaussian_count, 3]
+
+      raise ShapeError,
+            "expected quats [N,4] and scales [N,3], " \
+            "got #{data(quats).shape.inspect} and #{data(scales).shape.inspect}"
+    end
+    private_class_method :validate_geometry!
+
+    def validate_color_shape!(shape, gaussian_count, camera_count, sh_degree)
+      valid = if sh_degree
+                shape.length == 3 && shape[0] == gaussian_count && shape[1] >= ((sh_degree + 1)**2)
+              else
+                (shape.length == 2 && shape[0] == gaussian_count) ||
+                  (shape.length == 3 && shape[0...2] == [camera_count, gaussian_count])
+              end
+      return if valid
+
+      expected = sh_degree ? "[N,K,D] with K >= #{(sh_degree + 1)**2}" : "[N,D] or [C,N,D]"
+      raise ShapeError, "expected colors #{expected}, got #{shape.inspect}"
+    end
+    private_class_method :validate_color_shape!
+
+    def data(value)
+      Ops::TensorOps.data(value)
+    end
+    private_class_method :data
+  end
+end

data/lib/gsplat/strategy/base.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Strategy
+    # Shared strategy lifecycle and parameter/optimizer validation.
+    class Base
+      # Parameter names required by every structural strategy.
+      REQUIRED_KEYS = %i[means quats scales opacities sh0 shN].freeze
+
+      def initialize_state(scene_scale:)
+        raise ArgumentError, "scene_scale must be positive" unless scene_scale.positive?
+
+        { scene_scale: scene_scale }
+      end
+
+      # rubocop:disable Metrics/AbcSize, Naming/PredicateMethod
+      def check_sanity(params, optimizers)
+        missing_params = REQUIRED_KEYS - params.keys
+        missing_optimizers = REQUIRED_KEYS - optimizers.keys
+        raise ArgumentError, "missing params: #{missing_params.join(', ')}" unless missing_params.empty?
+        raise ArgumentError, "missing optimizers: #{missing_optimizers.join(', ')}" unless missing_optimizers.empty?
+
+        count = nil
+        REQUIRED_KEYS.each do |key|
+          variable = params.fetch(key)
+          raise ArgumentError, "params[:#{key}] must be an Autograd::Variable" unless variable.is_a?(Autograd::Variable)
+
+          count ||= variable.data.shape[0]
+          unless variable.data.ndim.positive? && variable.data.shape[0] == count
+            raise ShapeError, "all params must share first-axis size #{count}; #{key}=#{variable.data.shape.inspect}"
+          end
+
+          optimizer = optimizers.fetch(key)
+          unless optimizer.is_a?(Optim::Adam) &&
+                 optimizer.groups.values.any? { |group| group.variable.equal?(variable) }
+            raise ArgumentError, "optimizers[:#{key}] must optimize params[:#{key}]"
+          end
+        end
+        true
+      end
+      # rubocop:enable Metrics/AbcSize, Naming/PredicateMethod
+
+      def step_pre_backward(params:, optimizers:, state:, step:, info:)
+        [params, optimizers, state, step, info]
+        nil
+      end
+    end
+  end
+end