Add batch_size to ConstantField
@@ -10,12 +10,12 @@ from .field import Field, propagate_flags_children
@struct.definition()
class ConstantField(Field):
- def __init__(self, value=1.0, name=None, **kwargs):
+ def __init__(self, value=1.0, name=None, batch_size=1, **kwargs):
data = _convert_constant_to_data(value)
Field.__init__(self, **struct.kwargs(locals(), ignore='value'))
def sample_at(self, points):
- return _expand_axes(self.data, points)
+ return _expand_axes(self.data, points, self._batch_size)
@property
def rank(self):
@@ -50,13 +50,13 @@ def _convert_constant_to_data(value):
return value
-def _expand_axes(data, points):
+def _expand_axes(data, points, batch_size=1):
assert math.spatial_rank(data) >= 0
data = math.expand_dims(data, 1, math.spatial_rank(points) - math.spatial_rank(data))
points_axes = math.staticshape(points)[1:-1]
data_axes = math.staticshape(data)[1:-1]
for d_points, d_data in zip(points_axes, data_axes):
assert d_points % d_data == 0
- tilings = [1] + [d_points // d_data for d_points, d_data in zip(math.staticshape(points)[1:-1], math.staticshape(data)[1:-1])] + [1]
+ tilings = [batch_size or 1] + [d_points // d_data for d_points, d_data in zip(math.staticshape(points)[1:-1], math.staticshape(data)[1:-1])] + [1]
data = math.tile(data, tilings)
return data
