Class: Gsplat::Ops::ConcatFeatures

Inherits:
Autograd::Function show all
Defined in:
lib/gsplat/ops/tensor_shape_ops.rb

Overview

Concatenates any number of tensors along their final feature axis.

Class Method Summary collapse

Methods inherited from Autograd::Function

apply

Class Method Details

.backward(context, gradient) ⇒ Object



162
163
164
165
166
167
168
169
# File 'lib/gsplat/ops/tensor_shape_ops.rb', line 162

def backward(context, gradient)
  cursor = 0
  context.saved_values.first.map do |size|
    value = gradient[*Array.new(gradient.ndim - 1, true), cursor...(cursor + size)].dup
    cursor += size
    value
  end
end

.forward(context, *values) ⇒ Object

Raises:

  • (ArgumentError)


143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
# File 'lib/gsplat/ops/tensor_shape_ops.rb', line 143

def forward(context, *values)
  raise ArgumentError, "at least one feature tensor is required" if values.empty?

  leading_shape = values.first.shape[0...-1]
  unless values.all? { |value| value.shape[0...-1] == leading_shape }
    raise ShapeError, "feature tensors must share leading dimensions"
  end

  sizes = values.map { |value| value.shape[-1] }
  output = values.first.class.zeros(*(leading_shape + [sizes.sum]))
  cursor = 0
  values.each_with_index do |value, index|
    output[*Array.new(output.ndim - 1, true), cursor...(cursor + sizes[index])] = value
    cursor += sizes[index]
  end
  context.save(sizes)
  output
end