gsplat 1.0.0

data/lib/gsplat/math/camera_distortion.rb

@@ -0,0 +1,123 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Math
+    # OpenCV pinhole distortion and equidistant fisheye projection.
+    module CameraDistortion
+      module_function
+
+      # rubocop:disable Metrics/ParameterLists
+      def project_camera(means, intrinsics, camera_model, radial: nil, tangential: nil, thin_prism: nil)
+        # rubocop:enable Metrics/ParameterLists
+        projected = means.class.zeros(means.shape[0], 2)
+        jacobians = means.class.zeros(means.shape[0], 2, 3)
+        means.shape[0].times do |index|
+          point = means[index, true].to_a.map(&:to_f)
+          projected[index, true] = project_point(
+            point, intrinsics, camera_model, radial, tangential, thin_prism
+          )
+          jacobians[index, true, true] = numerical_jacobian(
+            point, intrinsics, camera_model, radial, tangential, thin_prism, means.class
+          )
+        end
+        [projected, jacobians]
+      end
+
+      # rubocop:disable Metrics/AbcSize, Metrics/ParameterLists
+      def project_point(point, intrinsics, camera_model, radial, tangential, thin_prism)
+        # rubocop:enable Metrics/AbcSize, Metrics/ParameterLists
+        x_coord = point[0] / point[2]
+        y_coord = point[1] / point[2]
+        normalized = if camera_model.to_s == "fisheye"
+                       fisheye_normalized = fisheye(x_coord, y_coord, radial)
+                       tangential_prism(
+                         *fisheye_normalized, *fisheye_normalized, tangential, thin_prism
+                       )
+                     else
+                       pinhole_normalized = pinhole(x_coord, y_coord, radial)
+                       tangential_prism(
+                         *pinhole_normalized, x_coord, y_coord, tangential, thin_prism
+                       )
+                     end
+        [
+          (intrinsics[0, 0].to_f * normalized[0]) +
+            (intrinsics[0, 1].to_f * normalized[1]) + intrinsics[0, 2].to_f,
+          (intrinsics[1, 0].to_f * normalized[0]) +
+            (intrinsics[1, 1].to_f * normalized[1]) + intrinsics[1, 2].to_f
+        ]
+      end
+
+      def pinhole(x_coord, y_coord, radial)
+        coefficients = Array.new(6, 0.0)
+        radial&.each_with_index { |value, index| coefficients[index] = value.to_f }
+        radius2 = (x_coord * x_coord) + (y_coord * y_coord)
+        radius4 = radius2 * radius2
+        radius6 = radius4 * radius2
+        numerator = 1 + (coefficients[0] * radius2) +
+                    (coefficients[1] * radius4) + (coefficients[2] * radius6)
+        denominator = 1 + (coefficients[3] * radius2) +
+                      (coefficients[4] * radius4) + (coefficients[5] * radius6)
+        scale = numerator / denominator
+        [x_coord * scale, y_coord * scale]
+      end
+      private_class_method :pinhole
+
+      def fisheye(x_coord, y_coord, radial)
+        radius = ::Math.sqrt((x_coord * x_coord) + (y_coord * y_coord))
+        return [x_coord, y_coord] if radius < 1e-12
+
+        theta = ::Math.atan(radius)
+        theta2 = theta * theta
+        coefficients = Array.new(4, 0.0)
+        radial&.each_with_index { |value, index| coefficients[index] = value.to_f }
+        polynomial = 1.0
+        power = theta2
+        coefficients.each do |coefficient|
+          polynomial += coefficient * power
+          power *= theta2
+        end
+        scale = theta * polynomial / radius
+        [x_coord * scale, y_coord * scale]
+      end
+      private_class_method :fisheye
+
+      # rubocop:disable Metrics/ParameterLists
+      def tangential_prism(base_x, base_y, x_coord, y_coord, tangential, thin_prism)
+        # rubocop:enable Metrics/ParameterLists
+        p1, p2 = tangential ? tangential.map(&:to_f) : [0.0, 0.0]
+        s1, s2, s3, s4 = thin_prism ? thin_prism.map(&:to_f) : [0.0, 0.0, 0.0, 0.0]
+        radius2 = (x_coord * x_coord) + (y_coord * y_coord)
+        radius4 = radius2 * radius2
+        [
+          base_x + (2 * p1 * x_coord * y_coord) + (p2 * (radius2 + (2 * x_coord * x_coord))) +
+            (s1 * radius2) + (s2 * radius4),
+          base_y + (p1 * (radius2 + (2 * y_coord * y_coord))) + (2 * p2 * x_coord * y_coord) +
+            (s3 * radius2) + (s4 * radius4)
+        ]
+      end
+      private_class_method :tangential_prism
+
+      # rubocop:disable Metrics/ParameterLists
+      def numerical_jacobian(point, intrinsics, camera_model, radial, tangential, thin_prism, type)
+        # rubocop:enable Metrics/ParameterLists
+        epsilon = type == Numo::DFloat ? 1e-6 : 1e-3
+        2.times.map do |output_axis|
+          3.times.map do |input_axis|
+            positive = point.dup
+            negative = point.dup
+            positive[input_axis] += epsilon
+            negative[input_axis] -= epsilon
+            projected_positive = project_point(
+              positive, intrinsics, camera_model, radial, tangential, thin_prism
+            )
+            projected_negative = project_point(
+              negative, intrinsics, camera_model, radial, tangential, thin_prism
+            )
+            (projected_positive[output_axis] - projected_negative[output_axis]) / (2 * epsilon)
+          end
+        end
+      end
+      private_class_method :numerical_jacobian
+    end
+  end
+end

data/lib/gsplat/math/camera_projection.rb

@@ -0,0 +1,202 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Math
+    # Camera-space transformation and pinhole Gaussian projection primitives.
+    module CameraProjection
+      module_function
+
+      def world_to_cam(means, covars, viewmats)
+        means, covars, viewmats = validate_world_inputs(means, covars, viewmats)
+        camera_count = viewmats.shape[0]
+        gaussian_count = means.shape[0]
+        camera_means = means.class.zeros(camera_count, gaussian_count, 3)
+        camera_covars = means.class.zeros(camera_count, gaussian_count, 3, 3)
+        camera_count.times do |camera_index|
+          rotation = viewmats[camera_index, 0...3, 0...3]
+          translation = viewmats[camera_index, 0...3, 3]
+          rotation_batch = means.class.zeros(gaussian_count, 3, 3)
+          rotation_batch[true, true, true] = rotation
+          camera_means[camera_index, true, true] = means.dot(rotation.transpose) + translation
+          rotated = Mat.matmul_batch(rotation_batch, covars)
+          camera_covars[camera_index, true, true, true] = Mat.matmul_batch(
+            rotated,
+            rotation_batch.transpose(0, 2, 1)
+          )
+        end
+        [camera_means, camera_covars]
+      end
+
+      def persp_proj(means, covars, intrinsics, width, height)
+        means, covars, intrinsics = validate_projection_inputs(means, covars, intrinsics, width, height)
+        camera_count = means.shape[0]
+        gaussian_count = means.shape[1]
+        means2d = means.class.zeros(camera_count, gaussian_count, 2)
+        covars2d = means.class.zeros(camera_count, gaussian_count, 2, 2)
+        camera_count.times do |camera_index|
+          projected_means, jacobians = pinhole_camera(
+            means[camera_index, true, true],
+            intrinsics[camera_index, true, true],
+            width,
+            height
+          )
+          means2d[camera_index, true, true] = projected_means
+          covariance = covars[camera_index, true, true, true]
+          transformed = Mat.matmul_batch(jacobians, covariance)
+          covars2d[camera_index, true, true, true] = Mat.matmul_batch(
+            transformed,
+            jacobians.transpose(0, 2, 1)
+          )
+        end
+        [means2d, covars2d]
+      end
+
+      def ortho_proj(means, covars, intrinsics, width, height)
+        means, covars, intrinsics = validate_projection_inputs(means, covars, intrinsics, width, height)
+        camera_count = means.shape[0]
+        gaussian_count = means.shape[1]
+        means2d = means.class.zeros(camera_count, gaussian_count, 2)
+        covars2d = means.class.zeros(camera_count, gaussian_count, 2, 2)
+        camera_count.times do |camera_index|
+          projected_means, jacobians = ortho_camera(
+            means[camera_index, true, true],
+            intrinsics[camera_index, true, true]
+          )
+          means2d[camera_index, true, true] = projected_means
+          transformed = Mat.matmul_batch(jacobians, covars[camera_index, true, true, true])
+          covars2d[camera_index, true, true, true] = Mat.matmul_batch(
+            transformed,
+            jacobians.transpose(0, 2, 1)
+          )
+        end
+        [means2d, covars2d]
+      end
+
+      # rubocop:disable Metrics/ParameterLists
+      def distorted_proj(means, covars, intrinsics, camera_model, radial_coeffs: nil,
+                         tangential_coeffs: nil, thin_prism_coeffs: nil)
+        # rubocop:enable Metrics/ParameterLists
+        camera_count = means.shape[0]
+        gaussian_count = means.shape[1]
+        means2d = means.class.zeros(camera_count, gaussian_count, 2)
+        covars2d = means.class.zeros(camera_count, gaussian_count, 2, 2)
+        camera_count.times do |camera_index|
+          projected, jacobians = CameraDistortion.project_camera(
+            means[camera_index, true, true],
+            intrinsics[camera_index, true, true],
+            camera_model,
+            radial: radial_coeffs && radial_coeffs[camera_index, true].to_a,
+            tangential: tangential_coeffs && tangential_coeffs[camera_index, true].to_a,
+            thin_prism: thin_prism_coeffs && thin_prism_coeffs[camera_index, true].to_a
+          )
+          means2d[camera_index, true, true] = projected
+          transformed = Mat.matmul_batch(jacobians, covars[camera_index, true, true, true])
+          covars2d[camera_index, true, true, true] = Mat.matmul_batch(
+            transformed,
+            jacobians.transpose(0, 2, 1)
+          )
+        end
+        [means2d, covars2d]
+      end
+
+      def validate_world_inputs(means, covars, viewmats)
+        validate_float_array!(means, "means")
+        covars = means.class.cast(covars)
+        viewmats = means.class.cast(viewmats)
+        unless means.ndim == 2 && means.shape[-1] == 3 && covars.shape == [means.shape[0], 3, 3]
+          raise ShapeError,
+                "expected means [N,3] and covars [N,3,3], " \
+                "got #{means.shape.inspect} and #{covars.shape.inspect}"
+        end
+        unless viewmats.ndim == 3 && viewmats.shape[1..] == [4, 4]
+          raise ShapeError, "expected viewmats [C,4,4], got #{viewmats.shape.inspect}"
+        end
+
+        [means, covars, viewmats]
+      end
+      private_class_method :validate_world_inputs
+
+      def validate_projection_inputs(means, covars, intrinsics, width, height)
+        validate_float_array!(means, "means")
+        covars = means.class.cast(covars)
+        intrinsics = means.class.cast(intrinsics)
+        valid = means.ndim == 3 && means.shape[-1] == 3 &&
+                covars.shape == means.shape[0...-1] + [3, 3] &&
+                intrinsics.shape == [means.shape[0], 3, 3]
+        unless valid
+          raise ShapeError,
+                "expected means [C,N,3], covars [C,N,3,3], and intrinsics [C,3,3], " \
+                "got #{means.shape.inspect}, #{covars.shape.inspect}, and #{intrinsics.shape.inspect}"
+        end
+        unless width.is_a?(Integer) && width.positive? && height.is_a?(Integer) && height.positive?
+          raise ArgumentError, "width and height must be positive integers"
+        end
+
+        [means, covars, intrinsics]
+      end
+      private_class_method :validate_projection_inputs
+
+      def validate_float_array!(array, name)
+        return if array.is_a?(Numo::NArray) && [Numo::SFloat, Numo::DFloat].include?(array.class)
+
+        raise ArgumentError, "#{name} must be Numo::SFloat or Numo::DFloat"
+      end
+      private_class_method :validate_float_array!
+
+      # rubocop:disable Metrics/AbcSize
+      def pinhole_camera(means, intrinsics, width, height)
+        x_coord = means[true, 0]
+        y_coord = means[true, 1]
+        z_coord = means[true, 2]
+        focal_x = intrinsics[0, 0]
+        focal_y = intrinsics[1, 1]
+        center_x = intrinsics[0, 2]
+        center_y = intrinsics[1, 2]
+        clipped_x = clip_ratio(
+          x_coord / z_coord,
+          -((center_x / focal_x) + (0.15 * width / focal_x)),
+          ((width - center_x) / focal_x) + (0.15 * width / focal_x)
+        ) * z_coord
+        clipped_y = clip_ratio(
+          y_coord / z_coord,
+          -((center_y / focal_y) + (0.15 * height / focal_y)),
+          ((height - center_y) / focal_y) + (0.15 * height / focal_y)
+        ) * z_coord
+        jacobians = means.class.zeros(means.shape[0], 2, 3)
+        jacobians[true, 0, 0] = focal_x / z_coord
+        jacobians[true, 0, 2] = -focal_x * clipped_x / (z_coord**2)
+        jacobians[true, 1, 1] = focal_y / z_coord
+        jacobians[true, 1, 2] = -focal_y * clipped_y / (z_coord**2)
+        homogeneous_x = (intrinsics[0, 0] * x_coord) + (intrinsics[0, 1] * y_coord) +
+                        (intrinsics[0, 2] * z_coord)
+        homogeneous_y = (intrinsics[1, 0] * x_coord) + (intrinsics[1, 1] * y_coord) +
+                        (intrinsics[1, 2] * z_coord)
+        means2d = means.class.zeros(means.shape[0], 2)
+        means2d[true, 0] = homogeneous_x / z_coord
+        means2d[true, 1] = homogeneous_y / z_coord
+        [means2d, jacobians]
+      end
+      # rubocop:enable Metrics/AbcSize
+
+      def clip_ratio(values, minimum, maximum)
+        output = values.dup
+        below = output.lt(minimum)
+        above = output.gt(maximum)
+        output[below] = minimum if below.any?
+        output[above] = maximum if above.any?
+        output
+      end
+      private_class_method :clip_ratio
+
+      def ortho_camera(means, intrinsics)
+        jacobians = means.class.zeros(means.shape[0], 2, 3)
+        jacobians[true, 0, 0] = intrinsics[0, 0]
+        jacobians[true, 1, 1] = intrinsics[1, 1]
+        means2d = means.class.zeros(means.shape[0], 2)
+        means2d[true, 0] = (means[true, 0] * intrinsics[0, 0]) + intrinsics[0, 2]
+        means2d[true, 1] = (means[true, 1] * intrinsics[1, 1]) + intrinsics[1, 2]
+        [means2d, jacobians]
+      end
+    end
+  end
+end

data/lib/gsplat/math/mat.rb

@@ -0,0 +1,114 @@
+# frozen_string_literal: true
+
+require_relative "small_matrix_primitives"
+
+module Gsplat
+  # Quaternion, projection, image metric, and small-matrix primitives.
+  module Math
+    # Closed-form batched operations for small matrices.
+    module Mat
+      extend SmallMatrixPrimitives
+
+      module_function
+
+      # Determinant of batched 2x2 matrices.
+      #
+      # @param matrices [Numo::NArray] [...,2,2]
+      # @param dtype [Class, Symbol, nil] optional float32/float64 calculation type
+      # @return [Numo::NArray] [...]
+      def det2x2(matrices, dtype: nil)
+        _, leading_shape, flat = prepare_matrix(matrices, 2, dtype)
+        determinant = (flat[true, 0, 0] * flat[true, 1, 1]) - (flat[true, 0, 1] * flat[true, 1, 0])
+        reshape_vector(determinant, leading_shape)
+      end
+
+      # Inverse of batched 2x2 matrices.
+      #
+      # @param matrices [Numo::NArray] [...,2,2]
+      # @param dtype [Class, Symbol, nil] optional float32/float64 calculation type
+      # @param epsilon [Numeric, nil] singularity threshold
+      # @return [Numo::NArray] [...,2,2]
+      def inv2x2(matrices, dtype: nil, epsilon: nil)
+        typed, leading_shape, flat = prepare_matrix(matrices, 2, dtype)
+        determinant = (flat[true, 0, 0] * flat[true, 1, 1]) - (flat[true, 0, 1] * flat[true, 1, 0])
+        ensure_invertible!(determinant, epsilon || default_epsilon(typed.class))
+        output = typed.class.zeros(flat.shape[0], 2, 2)
+        output[true, 0, 0] = flat[true, 1, 1] / determinant
+        output[true, 0, 1] = -flat[true, 0, 1] / determinant
+        output[true, 1, 0] = -flat[true, 1, 0] / determinant
+        output[true, 1, 1] = flat[true, 0, 0] / determinant
+        reshape_matrix(output, leading_shape, 2, 2)
+      end
+
+      # Ordered real eigenvalues of batched 2x2 matrices.
+      #
+      # @param matrices [Numo::NArray] [...,2,2]
+      # @param dtype [Class, Symbol, nil] optional float32/float64 calculation type
+      # @return [Numo::NArray] [...,2] in ascending order
+      def eigvals2x2(matrices, dtype: nil)
+        typed, leading_shape, flat = prepare_matrix(matrices, 2, dtype)
+        half_trace = (flat[true, 0, 0] + flat[true, 1, 1]) * 0.5
+        half_difference = (flat[true, 0, 0] - flat[true, 1, 1]) * 0.5
+        discriminant = (half_difference**2) + (flat[true, 0, 1] * flat[true, 1, 0])
+        discriminant[discriminant.lt(0)] = 0
+        root = discriminant**0.5
+        output = typed.class.zeros(flat.shape[0], 2)
+        output[true, 0] = half_trace - root
+        output[true, 1] = half_trace + root
+        reshape_matrix(output, leading_shape, 2)
+      end
+
+      # Determinant of batched 3x3 matrices.
+      #
+      # @param matrices [Numo::NArray] [...,3,3]
+      # @param dtype [Class, Symbol, nil] optional float32/float64 calculation type
+      # @return [Numo::NArray] [...]
+      def det3x3(matrices, dtype: nil)
+        _, leading_shape, flat = prepare_matrix(matrices, 3, dtype)
+        determinant = determinant3(flat)
+        reshape_vector(determinant, leading_shape)
+      end
+
+      # Inverse of batched 3x3 matrices.
+      #
+      # @param matrices [Numo::NArray] [...,3,3]
+      # @param dtype [Class, Symbol, nil] optional float32/float64 calculation type
+      # @param epsilon [Numeric, nil] singularity threshold
+      # @return [Numo::NArray] [...,3,3]
+      def inv3x3(matrices, dtype: nil, epsilon: nil)
+        typed, leading_shape, flat = prepare_matrix(matrices, 3, dtype)
+        determinant = determinant3(flat)
+        ensure_invertible!(determinant, epsilon || default_epsilon(typed.class))
+        output = typed.class.zeros(flat.shape[0], 3, 3)
+        fill_inverse3!(output, flat, determinant)
+        reshape_matrix(output, leading_shape, 3, 3)
+      end
+
+      # Batched matrix multiplication for small matrices.
+      #
+      # @param left [Numo::NArray] [...,M,K]
+      # @param right [Numo::NArray] [...,K,N]
+      # @param dtype [Class, Symbol, nil] optional float32/float64 calculation type
+      # @return [Numo::NArray] [...,M,N]
+      def matmul_batch(left, right, dtype: nil)
+        left = cast_float(left, dtype)
+        right = cast_float(right, left.class)
+        validate_rank!(left, 2)
+        validate_rank!(right, 2)
+        leading_shape = left.shape[0...-2]
+        unless right.shape[0...-2] == leading_shape && left.shape[-1] == right.shape[-2]
+          raise ShapeError, "matmul shape mismatch: left #{left.shape.inspect}, right #{right.shape.inspect}"
+        end
+
+        rows = left.shape[-2]
+        shared = left.shape[-1]
+        columns = right.shape[-1]
+        batch_size = leading_shape.empty? ? 1 : leading_shape.inject(:*)
+        left_flat = left.reshape(batch_size, rows, shared)
+        right_flat = right.reshape(batch_size, shared, columns)
+        output = multiply_flat(left_flat, right_flat, rows, shared, columns)
+        reshape_matrix(output, leading_shape, rows, columns)
+      end
+    end
+  end
+end

data/lib/gsplat/math/quaternion.rb

@@ -0,0 +1,175 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Math
+    # Batched wxyz quaternion normalization, rotation conversion, and VJPs.
+    module Quaternion
+      module_function
+
+      # Normalizes wxyz quaternions along the last dimension.
+      #
+      # @param quaternions [Numo::NArray] [...,4]
+      # @param dtype [Class, Symbol, nil] optional float32/float64 calculation type
+      # @param epsilon [Numeric, nil] minimum norm
+      # @return [Numo::NArray] [...,4]
+      def normalize(quaternions, dtype: nil, epsilon: nil)
+        typed, leading_shape, flat = prepare(quaternions, 4, dtype)
+        denominator = norm_denominator(flat, epsilon || default_epsilon(typed.class))
+        (flat / denominator.reshape(flat.shape[0], 1)).reshape(*(leading_shape + [4]))
+      end
+
+      # VJP of quaternion normalization.
+      #
+      # @param quaternions [Numo::NArray] [...,4]
+      # @param grad_output [Numo::NArray] [...,4]
+      # @param dtype [Class, Symbol, nil] optional float32/float64 calculation type
+      # @param epsilon [Numeric, nil] minimum norm
+      # @return [Numo::NArray] [...,4]
+      def normalize_vjp(quaternions, grad_output, dtype: nil, epsilon: nil)
+        typed, leading_shape, flat = prepare(quaternions, 4, dtype)
+        gradient, = prepare_matching(grad_output, leading_shape, 4, typed.class)
+        denominator = norm_denominator(flat, epsilon || default_epsilon(typed.class))
+        normalized = flat / denominator.reshape(flat.shape[0], 1)
+        dot = (normalized * gradient).sum(axis: 1)
+        result = (gradient - (normalized * dot.reshape(flat.shape[0], 1))) / denominator.reshape(flat.shape[0], 1)
+        result.reshape(*(leading_shape + [4]))
+      end
+
+      # Converts wxyz quaternions to rotation matrices.
+      #
+      # @param quaternions [Numo::NArray] [...,4]
+      # @param dtype [Class, Symbol, nil] optional float32/float64 calculation type
+      # @return [Numo::NArray] [...,3,3]
+      def to_rotmat(quaternions, dtype: nil)
+        typed, leading_shape, = prepare(quaternions, 4, dtype)
+        flat = normalize(typed).reshape(typed.size / 4, 4)
+        quat_w = flat[true, 0]
+        quat_x = flat[true, 1]
+        quat_y = flat[true, 2]
+        quat_z = flat[true, 3]
+        output = typed.class.zeros(flat.shape[0], 3, 3)
+        fill_rotation!(output, quat_w, quat_x, quat_y, quat_z)
+        output.reshape(*(leading_shape + [3, 3]))
+      end
+
+      # VJP of quaternion-to-rotation-matrix conversion.
+      #
+      # @param quaternions [Numo::NArray] [...,4]
+      # @param grad_output [Numo::NArray] [...,3,3]
+      # @param dtype [Class, Symbol, nil] optional float32/float64 calculation type
+      # @return [Numo::NArray] [...,4]
+      def to_rotmat_vjp(quaternions, grad_output, dtype: nil)
+        typed, leading_shape, = prepare(quaternions, 4, dtype)
+        gradient = cast_float(grad_output, typed.class)
+        unless gradient.shape == leading_shape + [3, 3]
+          raise ShapeError, "expected gradient #{(leading_shape + [3, 3]).inspect}, got #{gradient.shape.inspect}"
+        end
+
+        normalized = normalize(typed).reshape(typed.size / 4, 4)
+        grad_matrix = gradient.reshape(normalized.shape[0], 3, 3)
+        grad_normalized = rotation_vjp(normalized, grad_matrix)
+        normalize_vjp(typed, grad_normalized.reshape(*(leading_shape + [4])))
+      end
+
+      def prepare(quaternions, width, dtype)
+        typed = cast_float(quaternions, dtype)
+        unless typed.ndim >= 1 && typed.shape[-1] == width
+          raise ShapeError, "expected [...,#{width}], got #{typed.shape.inspect}"
+        end
+
+        leading_shape = typed.shape[0...-1]
+        [typed, leading_shape, typed.reshape(typed.size / width, width)]
+      end
+      private_class_method :prepare
+
+      def prepare_matching(value, leading_shape, width, dtype)
+        typed = cast_float(value, dtype)
+        expected = leading_shape + [width]
+        raise ShapeError, "expected #{expected.inspect}, got #{typed.shape.inspect}" unless typed.shape == expected
+
+        [typed.reshape(typed.size / width, width), leading_shape]
+      end
+      private_class_method :prepare_matching
+
+      def cast_float(array, dtype)
+        raise ArgumentError, "expected a Numo::NArray" unless array.is_a?(Numo::NArray)
+
+        type = case dtype
+               when nil then array.class
+               when :float32 then Numo::SFloat
+               when :float64 then Numo::DFloat
+               else dtype
+               end
+        unless [Numo::SFloat, Numo::DFloat].include?(type)
+          raise ArgumentError, "dtype must be Numo::SFloat or Numo::DFloat"
+        end
+
+        type.cast(array)
+      end
+      private_class_method :cast_float
+
+      def norm_denominator(flat, epsilon)
+        norms = (flat**2).sum(axis: 1)**0.5
+        norms[norms.lt(epsilon)] = epsilon
+        norms
+      end
+      private_class_method :norm_denominator
+
+      def default_epsilon(type)
+        type == Numo::DFloat ? 1e-12 : 1e-6
+      end
+      private_class_method :default_epsilon
+
+      # rubocop:disable Metrics/AbcSize
+      def fill_rotation!(output, quat_w, quat_x, quat_y, quat_z)
+        output[true, 0, 0] = 1 - (2 * ((quat_y**2) + (quat_z**2)))
+        output[true, 0, 1] = 2 * ((quat_x * quat_y) - (quat_w * quat_z))
+        output[true, 0, 2] = 2 * ((quat_x * quat_z) + (quat_w * quat_y))
+        output[true, 1, 0] = 2 * ((quat_x * quat_y) + (quat_w * quat_z))
+        output[true, 1, 1] = 1 - (2 * ((quat_x**2) + (quat_z**2)))
+        output[true, 1, 2] = 2 * ((quat_y * quat_z) - (quat_w * quat_x))
+        output[true, 2, 0] = 2 * ((quat_x * quat_z) - (quat_w * quat_y))
+        output[true, 2, 1] = 2 * ((quat_y * quat_z) + (quat_w * quat_x))
+        output[true, 2, 2] = 1 - (2 * ((quat_x**2) + (quat_y**2)))
+      end
+      # rubocop:enable Metrics/AbcSize
+      private_class_method :fill_rotation!
+
+      # rubocop:disable Metrics/AbcSize
+      def rotation_vjp(quaternion, gradient)
+        quat_w = quaternion[true, 0]
+        quat_x = quaternion[true, 1]
+        quat_y = quaternion[true, 2]
+        quat_z = quaternion[true, 3]
+        grad_at = ->(row, column) { gradient[true, row, column] }
+        output = quaternion.class.zeros(*quaternion.shape)
+        output[true, 0] = (
+          (-2 * quat_z * grad_at.call(0, 1)) + (2 * quat_y * grad_at.call(0, 2)) +
+          (2 * quat_z * grad_at.call(1, 0)) - (2 * quat_x * grad_at.call(1, 2)) -
+          (2 * quat_y * grad_at.call(2, 0)) + (2 * quat_x * grad_at.call(2, 1))
+        )
+        output[true, 1] = (
+          (2 * quat_y * grad_at.call(0, 1)) + (2 * quat_z * grad_at.call(0, 2)) +
+          (2 * quat_y * grad_at.call(1, 0)) - (4 * quat_x * grad_at.call(1, 1)) -
+          (2 * quat_w * grad_at.call(1, 2)) + (2 * quat_z * grad_at.call(2, 0)) +
+          (2 * quat_w * grad_at.call(2, 1)) - (4 * quat_x * grad_at.call(2, 2))
+        )
+        output[true, 2] = (
+          (-4 * quat_y * grad_at.call(0, 0)) + (2 * quat_x * grad_at.call(0, 1)) +
+          (2 * quat_w * grad_at.call(0, 2)) + (2 * quat_x * grad_at.call(1, 0)) +
+          (2 * quat_z * grad_at.call(1, 2)) - (2 * quat_w * grad_at.call(2, 0)) +
+          (2 * quat_z * grad_at.call(2, 1)) - (4 * quat_y * grad_at.call(2, 2))
+        )
+        output[true, 3] = (
+          (-4 * quat_z * grad_at.call(0, 0)) - (2 * quat_w * grad_at.call(0, 1)) +
+          (2 * quat_x * grad_at.call(0, 2)) + (2 * quat_w * grad_at.call(1, 0)) -
+          (4 * quat_z * grad_at.call(1, 1)) + (2 * quat_y * grad_at.call(1, 2)) +
+          (2 * quat_x * grad_at.call(2, 0)) + (2 * quat_y * grad_at.call(2, 1))
+        )
+        output
+      end
+      # rubocop:enable Metrics/AbcSize
+      private_class_method :rotation_vjp
+    end
+  end
+end