def create_uniform_pixel_coords_image(image_dims, batch_shape=None, normalized=False, dev_str=None):
"""
Create image of homogeneous integer :math:`xy` pixel co-ordinates :math:`\mathbf{X}\in\mathbb{Z}^{h×w×3}`, stored
as floating point values. The origin is at the top-left corner of the image, with :math:`+x` rightwards, and
:math:`+y` downwards. The final homogeneous dimension are all ones. In subsequent use of this image, the depth of
each pixel can be represented using this same homogeneous representation, by simply scaling each 3-vector by the
depth value. The final dimension therefore always holds the depth value, while the former two dimensions hold depth
scaled pixel co-ordinates.\n
`[reference] <localhost:63342/ivy/docs/source/references/mvg_textbook.pdf#page=172>`_
deduction from top of page 154, section 6.1, equation 6.1
:param image_dims: Image dimensions.
:type image_dims: sequence of ints.
:param batch_shape: Shape of batch. Assumed no batch dimensions if None.
:type batch_shape: sequence of ints, optional
:param normalized: Whether to normalize x-y pixel co-ordinates to the range 0-1.
:type normalized: bool
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc.
:type dev_str: str, optional
:return: Image of homogeneous pixel co-ordinates *[batch_shape,height,width,3]*
"""
# shapes as lists
batch_shape = [] if batch_shape is None else batch_shape
batch_shape = list(batch_shape)
image_dims = list(image_dims)
# other shape specs
num_batch_dims = len(batch_shape)
flat_shape = [1] * num_batch_dims + image_dims + [3]
tile_shape = batch_shape + [1] * 3
# H x W x 1
pixel_x_coords = _ivy.cast(_ivy.reshape(_ivy.tile(_ivy.arange(image_dims[1], dev_str=dev_str), [image_dims[0]]),
(image_dims[0], image_dims[1], 1)), 'float32')
if normalized:
pixel_x_coords = pixel_x_coords / (float(image_dims[1]) + MIN_DENOMINATOR)
# W x H x 1
pixel_y_coords_ = _ivy.cast(_ivy.reshape(_ivy.tile(_ivy.arange(image_dims[0], dev_str=dev_str), [image_dims[1]]),
(image_dims[1], image_dims[0], 1)), 'float32')
# H x W x 1
pixel_y_coords = _ivy.transpose(pixel_y_coords_, (1, 0, 2))
if normalized:
pixel_y_coords = pixel_y_coords / (float(image_dims[0]) + MIN_DENOMINATOR)
# H x W x 1
ones = _ivy.ones_like(pixel_x_coords, dev_str=dev_str)
# BS x H x W x 3
return _ivy.tile(_ivy.reshape(_ivy.concatenate((pixel_x_coords, pixel_y_coords, ones), -1),
flat_shape), tile_shape)
def __init__(self,
img_meas: Dict[str, ESMCamMeasurement],
agent_rel_mat: ivy.Array,
control_mean: ivy.Array = None,
control_cov: ivy.Array = None):
"""
Create esm observation container
:param img_meas: dict of ESMImageMeasurement objects, with keys for camera names.
:type: img_meas: Ivy container
:param agent_rel_mat: The pose of the agent relative to the previous pose, in matrix form
*[batch_size, timesteps, 3, 4]*.
:type agent_rel_mat: array
:param control_mean: The pose of the agent relative to the previous pose, in rotation vector pose form.
Inferred from agent_rel_mat if None. *[batch_size, timesteps, 6]*
:type control_mean: array, optional
:param control_cov: The convariance of the agent relative pose, in rotation vector form.
Assumed all zero if None. *[batch_size, timesteps, 6, 6]*.
:type control_cov: array, optional
"""
self['img_meas'] = Container(img_meas)
agent_rel_mat = _pad_to_batch_n_time_dims(agent_rel_mat, 4)
self['agent_rel_mat'] = agent_rel_mat
if control_mean is None:
control_mean = ivy_mech.mat_pose_to_rot_vec_pose(agent_rel_mat)
else:
control_mean = _pad_to_batch_n_time_dims(control_mean, 3)
self['control_mean'] = control_mean
if control_cov is None:
control_cov = ivy.tile(ivy.expand_dims(ivy.zeros_like(control_mean), -1), (1, 1, 1, 6))
else:
control_cov = _pad_to_batch_n_time_dims(control_cov, 4)
self['control_cov'] = control_cov
def make_transformation_homogeneous(matrices, batch_shape=None, dev_str=None):
"""
Append to set of 3x4 non-homogeneous matrices to make them homogeneous.
:param matrices: set of 3x4 non-homogeneous matrices *[batch_shape,3,4]*
:type matrices: array
:param batch_shape: Shape of batch. Inferred from inputs if None.
:type batch_shape: sequence of ints, optional
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Same as x if None.
:type dev_str: str, optional
:return: 4x4 Homogeneous matrices *[batch_shape,4,4]*
"""
if batch_shape is None:
batch_shape = matrices.shape[:-2]
if dev_str is None:
dev_str = _ivy.dev_str(matrices)
# shapes as list
batch_shape = list(batch_shape)
num_batch_dims = len(batch_shape)
# BS x 1 x 4
last_row = _ivy.tile(
_ivy.reshape(_ivy.array([0., 0., 0., 1.], dev_str=dev_str),
[1] * num_batch_dims + [1, 4]), batch_shape + [1, 1])
# BS x 4 x 4
return _ivy.concatenate((matrices, last_row), -2)
def rot_mat_and_cam_center_to_ext_mat(rotation_mat, camera_center, batch_shape=None):
"""
Get extrinsic matrix :math:`\mathbf{E}\in\mathbb{R}^{3×4}` from rotation matrix
:math:`\mathbf{R}\in\mathbb{R}^{3×3}` and camera centers :math:`\overset{\sim}{\mathbf{C}}\in\mathbb{R}^{3×1}`.\n
`[reference] <localhost:63342/ivy/docs/source/references/mvg_textbook.pdf#page=175>`_
page 157, section 6.1, equation 6.11
:param rotation_mat: Rotation matrix *[batch_shape,3,3]*
:type rotation_mat: array
:param camera_center: Camera center *[batch_shape,3,1]*
:type camera_center: array
:param batch_shape: Shape of batch. Inferred from inputs if None.
:type batch_shape: sequence of ints, optional
:return: Extrinsic matrix *[batch_shape,3,4]*
"""
if batch_shape is None:
batch_shape = rotation_mat.shape[:-2]
# shapes as list
batch_shape = list(batch_shape)
# num batch dims
num_batch_dims = len(batch_shape)
# BS x 3 x 3
identity = _ivy.tile(_ivy.reshape(_ivy.identity(3), [1] * num_batch_dims + [3, 3]),
batch_shape + [1, 1])
# BS x 3 x 4
identity_w_cam_center = _ivy.concatenate((identity, -camera_center), -1)
# BS x 3 x 4
return _ivy.matmul(rotation_mat, identity_w_cam_center)
def __init__(self, num_iters, compile_flag, interactive, dev_str, f):
# ivy
f = choose_random_framework() if f is None else f
ivy.set_framework(f)
ivy.seed(0)
# device
if dev_str is None:
dev_str = 'gpu:0' if ivy.gpu_is_available() else 'cpu'
self._dev_str = dev_str
# Load input images and poses
this_dir = os.path.dirname(os.path.realpath(__file__))
data = np.load(os.path.join(this_dir, 'nerf_data/tiny_nerf_data.npz'))
images = ivy.array(data['images'], 'float32', dev_str)
inv_ext_mats = ivy.array(data['poses'], 'float32', dev_str)
# intrinsics
focal_lengths = ivy.array(np.tile(np.reshape(data['focal'], (1, 1)), [100, 2]), 'float32', dev_str)
self._img_dims = images.shape[1:3]
pp_offsets = ivy.tile(ivy.array([[dim/2 - 0.5 for dim in self._img_dims]], dev_str=dev_str), [100, 1])
# train data
self._images = images[:100, ..., :3]
self._intrinsics = ivy_vision.focal_lengths_and_pp_offsets_to_intrinsics_object(
focal_lengths, pp_offsets, self._img_dims)
self._cam_geoms = ivy_vision.inv_ext_mat_and_intrinsics_to_cam_geometry_object(
inv_ext_mats[:100, 0:3], self._intrinsics)
# test data
self._test_img = images[101]
self._test_cam_geom = ivy_vision.inv_ext_mat_and_intrinsics_to_cam_geometry_object(
inv_ext_mats[101, 0:3], self._intrinsics.slice(0))
# train config
self._embed_length = 6
self._lr = 5e-4
self._num_samples = 64
self._num_iters = num_iters
# log config
self._interactive = interactive
self._log_freq = 1
self._vis_freq = 25 if self._interactive else -1
self._vis_log_dir = 'nerf_renderings'
if os.path.exists(self._vis_log_dir):
shutil.rmtree(self._vis_log_dir)
os.makedirs(self._vis_log_dir)
# model
self._model = Model(4, 256, self._embed_length, dev_str)
# compile
if compile_flag:
rays_o, rays_d = self._get_rays(self._cam_geoms.slice(0))
target = self._images[0]
self._loss_fn = ivy.compile_fn(self._loss_fn, False,
example_inputs=[self._model, rays_o, rays_d, target, self._model.v])
def __init__(self,
img_mean: ivy.Array,
cam_rel_mat: ivy.Array,
img_var: ivy.Array = None,
validity_mask: ivy.Array = None,
pose_mean: ivy.Array = None,
pose_cov: ivy.Array = None):
"""
Create esm image measurement container
:param img_mean: Camera-relative co-ordinates and image features
*[batch_size, timesteps, height, width, 3 + feat]*
:type: img_mean: array
:param cam_rel_mat: The pose of the camera relative to the current agent pose, in matrix form
*[batch_size, timesteps, 3, 4]*
:type cam_rel_mat: array
:param img_var: Image depth and feature variance values, assumed all zero if None.
*[batch_size, timesteps, height, width, 1 + feat]*
:type: img_var: array, optional
:param validity_mask: Validity mask, for which pixels should be considered. Assumed all valid if None
*[batch_size, timesteps, height, width, 1]*
:type validity_mask: array, optional
:param pose_mean: The pose of the camera relative to the current agent pose, in rotation vector pose form.
Inferred from cam_rel_mat if None. *[batch_size, timesteps, 6]*
:type pose_mean: array, optional
:param pose_cov: The convariance of the camera relative pose, in rotation vector form. Assumed all zero if None.
*[batch_size, timesteps, 6, 6]*
:type pose_cov: array, optional
"""
img_mean = _pad_to_batch_n_time_dims(img_mean, 5)
cam_rel_mat = _pad_to_batch_n_time_dims(cam_rel_mat, 4)
self['img_mean'] = img_mean
self['cam_rel_mat'] = cam_rel_mat
if img_var is None:
img_var = ivy.zeros_like(img_mean)
else:
img_var = _pad_to_batch_n_time_dims(img_var, 5)
self['img_var'] = img_var
if validity_mask is None:
validity_mask = ivy.ones_like(img_mean[..., 0:1])
else:
validity_mask = _pad_to_batch_n_time_dims(validity_mask, 5)
self['validity_mask'] = validity_mask
if pose_mean is None:
pose_mean = ivy_mech.mat_pose_to_rot_vec_pose(cam_rel_mat)
else:
pose_mean = _pad_to_batch_n_time_dims(pose_mean, 3)
self['pose_mean'] = pose_mean
if pose_cov is None:
pose_cov = ivy.tile(ivy.expand_dims(ivy.zeros_like(pose_mean), -1), (1, 1, 1, 6))
else:
pose_cov = _pad_to_batch_n_time_dims(pose_cov, 4)
self['pose_cov'] = pose_cov
def create_trimesh_indices_for_image(batch_shape, image_dims, dev_str='cpu:0'):
"""
Create triangle mesh for image with given image dimensions
:param batch_shape: Shape of batch.
:type batch_shape: sequence of ints
:param image_dims: Image dimensions.
:type image_dims: sequence of ints
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc.
:type dev_str: str, optional
:return: Triangle mesh indices for image *[batch_shape,h*w*some_other_stuff,3]*
"""
# shapes as lists
batch_shape = list(batch_shape)
image_dims = list(image_dims)
# other shape specs
num_batch_dims = len(batch_shape)
tri_dim = 2 * (image_dims[0] - 1) * (image_dims[1] - 1)
flat_shape = [1] * num_batch_dims + [tri_dim] + [3]
tile_shape = batch_shape + [1] * 2
# 1 x W-1
t00_ = _ivy.reshape(_ivy.arange(image_dims[1] - 1, dtype_str='float32', dev_str=dev_str), (1, -1))
# H-1 x 1
k_ = _ivy.reshape(_ivy.arange(image_dims[0] - 1, dtype_str='float32', dev_str=dev_str), (-1, 1)) * image_dims[1]
# H-1 x W-1
t00_ = _ivy.matmul(_ivy.ones((image_dims[0] - 1, 1), dev_str=dev_str), t00_)
k_ = _ivy.matmul(k_, _ivy.ones((1, image_dims[1] - 1), dev_str=dev_str))
# (H-1xW-1) x 1
t00 = _ivy.expand_dims(t00_ + k_, -1)
t01 = t00 + 1
t02 = t00 + image_dims[1]
t10 = t00 + image_dims[1] + 1
t11 = t01
t12 = t02
# (H-1xW-1) x 3
t0 = _ivy.concatenate((t00, t01, t02), -1)
t1 = _ivy.concatenate((t10, t11, t12), -1)
# BS x 2x(H-1xW-1) x 3
return _ivy.tile(_ivy.reshape(_ivy.concatenate((t0, t1), 0),
flat_shape), tile_shape)
def _group_tensor_into_windowed_tensor_simple(self, x, seq_info):
seq_info = self._update_seq_info_for_window(seq_info)
if self._fixed_sequence_length:
return ivy.reshape(ivy.gather_nd(x, ivy.array(self._gather_idxs)),
(self._windows_per_seq, self._window_size) +
x.shape[1:])
else:
num_windows_in_seq = int(
ivy.to_numpy(
ivy.maximum(seq_info.length[0] - self._window_size + 1,
1)))
window_idxs_in_seq = ivy.arange(num_windows_in_seq, 0, 1)
gather_idxs = ivy.tile(
ivy.reshape(ivy.arange(self._window_size, 0, 1),
(1, self._window_size)),
(num_windows_in_seq, 1)) + ivy.expand_dims(
window_idxs_in_seq, -1)
gather_idxs_flat = ivy.reshape(
gather_idxs, (self._window_size * num_windows_in_seq, 1))
return ivy.reshape(ivy.gather_nd(x, gather_idxs_flat),
(num_windows_in_seq, self._window_size) +
x.shape[1:])
def smooth_image_fom_var_image(mean,
var,
kernel_dim,
kernel_scale,
dev_str=None):
"""
Smooth an image using variance values from a variance image of the same size, and a spatial smoothing kernel.
:param mean: Image to smooth *[batch_shape,h,w,d]*
:type mean: array
:param var: Variance image, with the variance values of each pixel in the image *[batch_shape,h,w,d]*
:type var: array
:param kernel_dim: The dimension of the kernel
:type kernel_dim: int
:param kernel_scale: The scale of the kernel along the channel dimension *[d]*
:type kernel_scale: array
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Same as x if None.
:type dev_str: str, optional
:return: Image smoothed based on variance image and smoothing kernel.
"""
if dev_str is None:
dev_str = _ivy.dev_str(mean)
# shapes as list
kernel_shape = [kernel_dim, kernel_dim]
kernel_size = kernel_dim**2
dims = mean.shape[-1]
# KH x KW x 2
uniform_pixel_coords = _ivy_svg.create_uniform_pixel_coords_image(
kernel_shape, dev_str=dev_str)[..., 0:2]
# 2
kernel_central_pixel_coord = _ivy.array([
float(_math.floor(kernel_shape[0] / 2)),
float(_math.floor(kernel_shape[1] / 2))
],
dev_str=dev_str)
# KH x KW x 2
kernel_xy_dists = kernel_central_pixel_coord - uniform_pixel_coords
kernel_xy_dists_sqrd = kernel_xy_dists**2
# KW x KW x D x D
unit_kernel = _ivy.tile(
_ivy.reduce_sum(kernel_xy_dists_sqrd, -1, keepdims=True)**0.5,
(1, 1, dims))
kernel = 1 + unit_kernel * kernel_scale
recip_kernel = 1 / (kernel + MIN_DENOMINATOR)
# D
kernel_sum = _ivy.reduce_sum(kernel, [0, 1])[0]
recip_kernel_sum = _ivy.reduce_sum(recip_kernel, [0, 1])
# BS x H x W x D
recip_var = 1 / (var + MIN_DENOMINATOR)
recip_var_scaled = recip_var + 1
recip_new_var_scaled = _ivy.depthwise_conv2d(recip_var_scaled,
recip_kernel, 1, "VALID")
# This 0.99 prevents float32 rounding errors leading to -ve variances, the true equation would use 1.0
recip_new_var = recip_new_var_scaled - recip_kernel_sum * 0.99
new_var = 1 / (recip_new_var + MIN_DENOMINATOR)
mean_x_recip_var = mean * recip_var
mean_x_recip_var_sum = _ivy.abs(
_ivy.depthwise_conv2d(mean_x_recip_var, recip_kernel, 1, "VALID"))
new_mean = new_var * mean_x_recip_var_sum
new_var = new_var * kernel_size**2 / (kernel_sum + MIN_DENOMINATOR)
# prevent overconfidence from false meas independence assumption
# BS x H x W x D, # BS x H x W x D
return new_mean, new_var
def _forward(self,
obs: ESMObservation,
memory: ESMMemory = None,
batch_size=None,
num_timesteps=None,
num_cams=None,
image_dims=None):
"""
Perform ESM update step.
:param obs: Observations
:type obs: ESMObservation
:param memory: Memory from the previous time-step, uses internal parameter if None.
:type memory: ESMMemory, optional.
:param batch_size: Size of batch, inferred from inputs if None.
:type batch_size: int, optional
:param num_timesteps: Number of timesteps, inferred from inputs if None.
:type num_timesteps: int, optional
:param num_cams: Number of cameras, inferred from inputs if None.
:type num_cams: int, optional
:param image_dims: Image dimensions of captured images, inferred from inputs if None.
:type image_dims: sequence of ints, optional
:return: New memory of type ESMMemory
"""
# get shapes
img_meas = (next(iter(obs.img_meas.values()))).img_mean
if batch_size is None:
batch_size = int(img_meas.shape[0])
if num_timesteps is None:
num_timesteps = int(img_meas.shape[1])
if num_cams is None:
num_cams = len(obs.img_meas.values())
if image_dims is None:
image_dims = list(img_meas.shape[2:4])
# get only previous memory
# extract from memory #
# --------------------#
if memory:
prev_mem = memory.slice((slice(None, None, None), -1))
elif self._stateful and self._memory is not None:
prev_mem = self._memory.slice((slice(None, None, None), -1))
else:
prev_mem = self.empty_memory(batch_size, 1).slice(
(slice(None, None, None), -1))
# holes prior #
# ------------#
# B x N x OH x OW x 1
self._sphere_depth_prior = \
ivy.ones([batch_size, num_timesteps] + self._sphere_img_dims + [1], dev_str=self._dev_str) * \
self._sphere_depth_prior_val
# B x N x OH x OW x F
self._sphere_feat_prior = \
ivy.ones([batch_size, num_timesteps] + self._sphere_img_dims + [self._feat_dim], dev_str=self._dev_str) * \
self._sphere_feat_prior_val
# B x N x OH x OW x (1+F)
holes_prior = ivy.concatenate([self._sphere_depth_prior] +
[self._sphere_feat_prior], -1)
# holes prior variances #
# ----------------------#
# B x N x OH x OW x 2
sphere_ang_pix_prior_var = \
ivy.ones([batch_size, num_timesteps] + self._sphere_img_dims + [2], dev_str=self._dev_str) * self._ang_pix_prior_var_val
# B x N x OH x OW x 1
sphere_depth_prior_var = \
ivy.ones([batch_size, num_timesteps] + self._sphere_img_dims + [1], dev_str=self._dev_str) * self._depth_prior_var_val
# B x N x OH x OW x F
sphere_feat_prior_var = \
ivy.ones([batch_size, num_timesteps] + self._sphere_img_dims + [self._feat_dim], dev_str=self._dev_str) * \
self._feat_prior_var_val
# B x N x OH x OW x (3+F)
holes_prior_var = ivy.concatenate([sphere_ang_pix_prior_var] +
[sphere_depth_prior_var] +
[sphere_feat_prior_var], -1)
# variance threshold #
# -------------------#
# B x N x (3+F) x 1
var_threshold_min = ivy.tile(
ivy.reshape(
ivy.stack([self._min_ang_pix_var] * 2 + [self._min_depth_var] +
[self._min_feat_var] * self._feat_dim),
[1, 1, 3 + self._feat_dim, 1]),
[batch_size, num_timesteps, 1, 1])
var_threshold_max = ivy.tile(
ivy.reshape(
ivy.stack([self._ang_pix_var_threshold] * 2 +
[self._depth_var_threshold] +
[self._feat_var_threshold] * self._feat_dim),
[1, 1, 3 + self._feat_dim, 1]),
[batch_size, num_timesteps, 1, 1])
self._var_threshold = ivy.concatenate(
(var_threshold_min, var_threshold_max), -1)
# measurements #
# -------------#
# B x N x OH x OW x 3
uniform_sphere_pixel_coords = ivy_vision.create_uniform_pixel_coords_image(
self._sphere_img_dims, (batch_size, num_timesteps),
dev_str=self._dev_str)
# B x N x OH x OW x (3+F), B x N x OH x OW x (3+F)
meas_means, meas_vars = self._convert_images_to_omni_observations(
obs.img_meas, uniform_sphere_pixel_coords, holes_prior, batch_size,
num_timesteps, num_cams, image_dims)
# filtering #
# ----------#
# list of B x OH x OW x (3+F), list of B x OH x OW x (3+F)
fused_measurements_list, fused_variances_list = \
self._kalman_filter_on_measurement_sequence(
prev_mem.mean, prev_mem.var, holes_prior[:, 0], holes_prior_var[:, 0], meas_means, meas_vars,
uniform_sphere_pixel_coords[:, 0], obs.control_mean, obs.control_cov, obs.agent_rel_mat,
batch_size, num_timesteps)
# new variance #
# -------------#
# B x N x OH x OW x (3+F)
fused_variance = ivy.concatenate(
[ivy.expand_dims(item, 1) for item in fused_variances_list], 1)
# variance clipping
# B x N x OH x OW x 1
fused_depth_variance = ivy.clip(fused_variance[..., 0:1],
self._min_depth_var,
self._depth_prior_var_val)
# B x N x OH x OW x 3
fused_feat_variance = ivy.clip(fused_variance[...,
1:], self._min_feat_var,
self._feat_prior_var_val)
# B x N x OH x OW x (3+F)
fused_variance = ivy.concatenate([fused_depth_variance] +
[fused_feat_variance], -1)
# new mean #
# ---------#
# B x N x OH x OW x (3+F)
fused_measurement = ivy.concatenate(
[ivy.expand_dims(item, 1) for item in fused_measurements_list], 1)
# value clipping
# B x N x OH x OW x 2
fused_pixel_coords = fused_measurement[..., 0:2]
# B x N x OH x OW x 1
fused_depth = ivy.clip(fused_measurement[..., 2:3], self._min_depth,
self._max_depth)
# B x N x OH x OW x 3
fused_feat = fused_measurement[..., 3:]
# B x N x OH x OW x (3+F)
fused_measurement = ivy.concatenate([fused_pixel_coords] +
[fused_depth] + [fused_feat], -1)
# update memory #
# --------------#
# B x N x OH x OW x (3+F), B x N x OH x OW x (3+F)
self._memory = ESMMemory(mean=fused_measurement, var=fused_variance)
# return #
# -------#
return self._memory
def _forward(self, x, prev_state):
prev_read_vector_list = prev_state[1]
controller_input = ivy.concatenate([x] + prev_read_vector_list, axis=1)
controller_output, controller_state = self._controller(ivy.expand_dims(controller_input, -2),
initial_state=prev_state[0])
controller_output = controller_output[..., -1, :]
parameters = self._controller_proj(controller_output)
parameters = ivy.clip(parameters, -self._clip_value, self._clip_value)
head_parameter_list = \
ivy.split(parameters[:, :self._num_parameters_per_head * self._num_heads], self._num_heads,
axis=1)
erase_add_list = ivy.split(parameters[:, self._num_parameters_per_head * self._num_heads:],
2 * self._write_head_num, axis=1)
prev_w_list = prev_state[2]
prev_M = prev_state[4]
w_list = []
for i, head_parameter in enumerate(head_parameter_list):
k = ivy.tanh(head_parameter[:, 0:self._memory_vector_dim])
beta = ivy.softplus(head_parameter[:, self._memory_vector_dim])
g = ivy.sigmoid(head_parameter[:, self._memory_vector_dim + 1])
s = ivy.softmax(
head_parameter[:, self._memory_vector_dim + 2:self._memory_vector_dim +
2 + (self._shift_range * 2 + 1)])
gamma = ivy.softplus(head_parameter[:, -1]) + 1
w = self._addressing(k, beta, g, s, gamma, prev_M, prev_w_list[i])
w_list.append(w)
# Reading (Sec 3.1)
read_w_list = w_list[:self._read_head_num]
if self._step == 0:
usage_indicator = ivy.zeros_like(w_list[0])
else:
usage_indicator = prev_state[3] + ivy.reduce_sum(ivy.concatenate(read_w_list, 0))
read_vector_list = []
for i in range(self._read_head_num):
read_vector = ivy.reduce_sum(ivy.expand_dims(read_w_list[i], axis=2) * prev_M, axis=1)
read_vector_list.append(read_vector)
# Writing (Sec 3.2)
prev_wrtie_w_list = prev_w_list[self._read_head_num:]
w_wr_size = math.ceil(self._memory_size / 2) if self._retroactive_updates else self._memory_size
if self._sequential_writing:
batch_size = ivy.shape(x)[0]
if self._step < w_wr_size:
w_wr_list = [ivy.tile(ivy.cast(ivy.one_hot(
ivy.array([self._step]), w_wr_size), 'float32'),
(batch_size, 1))] * self._write_head_num
else:
batch_idxs = ivy.expand_dims(ivy.arange(batch_size, 0), -1)
mem_idxs = ivy.expand_dims(ivy.argmax(usage_indicator[..., :w_wr_size], -1), -1)
total_idxs = ivy.concatenate((batch_idxs, mem_idxs), -1)
w_wr_list = [ivy.scatter_nd(total_idxs, ivy.ones((batch_size,)),
(batch_size, w_wr_size))] * self._write_head_num
else:
w_wr_list = w_list[self._read_head_num:]
if self._retroactive_updates:
w_ret_list = [self._retroactive_discount * prev_wrtie_w[..., w_wr_size:] +
(1 - self._retroactive_discount) * prev_wrtie_w[..., :w_wr_size]
for prev_wrtie_w in prev_wrtie_w_list]
w_wrtie_list = [ivy.concatenate((w_wr, w_ret), -1) for w_wr, w_ret in zip(w_wr_list, w_ret_list)]
else:
w_wrtie_list = w_wr_list
M = prev_M
for i in range(self._write_head_num):
w = ivy.expand_dims(w_wrtie_list[i], axis=2)
if self._with_erase:
erase_vector = ivy.expand_dims(ivy.sigmoid(erase_add_list[i * 2]), axis=1)
M = M * ivy.ones(ivy.shape(M)) - ivy.matmul(w, erase_vector)
add_vector = ivy.expand_dims(ivy.tanh(erase_add_list[i * 2 + 1]), axis=1)
M = M + ivy.matmul(w, add_vector)
NTM_output = self._output_proj(ivy.concatenate([controller_output] + read_vector_list, axis=1))
NTM_output = ivy.clip(NTM_output, -self._clip_value, self._clip_value)
self._step += 1
return NTM_output, NTMControllerState(
controller_state=controller_state, read_vector_list=read_vector_list, w_list=w_list,
usage_indicator=usage_indicator, M=M)
def _triangulate_depth_by_homogeneous_dlt(ds_pixel_coords, full_mats, _, _1,
batch_shape, image_dims):
# num batch dims
num_batch_dims = len(batch_shape)
# BS x 2 x H x W x 3
pixel_coords_normalized = ds_pixel_coords / (ds_pixel_coords[..., -1:] +
MIN_DENOMINATOR)
# BS x 3 x 4
P = full_mats[..., 0, :, :]
P_dash = full_mats[..., 1, :, :]
# BS x (HxW) x 4
p1T = _ivy.tile(P[..., 0:1, :],
[1] * num_batch_dims + [image_dims[0] * image_dims[1], 1])
p2T = _ivy.tile(P[..., 1:2, :],
[1] * num_batch_dims + [image_dims[0] * image_dims[1], 1])
p3T = _ivy.tile(P[..., 2:3, :],
[1] * num_batch_dims + [image_dims[0] * image_dims[1], 1])
p_dash_1T = _ivy.tile(P_dash[..., 0:1, :], [1] * num_batch_dims +
[image_dims[0] * image_dims[1], 1])
p_dash_2T = _ivy.tile(P_dash[..., 1:2, :], [1] * num_batch_dims +
[image_dims[0] * image_dims[1], 1])
p_dash_3T = _ivy.tile(P_dash[..., 2:3, :], [1] * num_batch_dims +
[image_dims[0] * image_dims[1], 1])
# BS x (WxH) x 1
x = _ivy.reshape(pixel_coords_normalized[..., 0, :, :, 0],
batch_shape + [-1, 1])
y = _ivy.reshape(pixel_coords_normalized[..., 0, :, :, 1],
batch_shape + [-1, 1])
x_dash = _ivy.reshape(pixel_coords_normalized[..., 1, :, :, 0],
batch_shape + [-1, 1])
y_dash = _ivy.reshape(pixel_coords_normalized[..., 1, :, :, 1],
batch_shape + [-1, 1])
# BS x (HxW) x 1 x 4
A_row1 = _ivy.expand_dims(x * p3T - p1T, -2)
A_row2 = _ivy.expand_dims(y * p3T - p2T, -2)
A_row3 = _ivy.expand_dims(x_dash * p_dash_3T - p_dash_1T, -2)
A_row4 = _ivy.expand_dims(y_dash * p_dash_3T - p_dash_2T, -2)
# BS x (HxW) x 4 x 4
A = _ivy.concatenate((A_row1, A_row2, A_row3, A_row4), -2)
# BS x (HxW) x 4
X = _ivy_pg.solve_homogeneous_dlt(A)
# BS x W x H x 4
coords_wrt_world_homo_unscaled = _ivy.reshape(
X, batch_shape + image_dims + [4])
coords_wrt_world = coords_wrt_world_homo_unscaled / (
coords_wrt_world_homo_unscaled[..., -1:] + MIN_DENOMINATOR)
# BS x W x H x 3
return _ivy_svg.world_to_ds_pixel_coords(coords_wrt_world,
full_mats[..., 0, :, :],
batch_shape, image_dims)
def coords_to_voxel_grid(coords,
voxel_shape_spec,
mode='DIMS',
coord_bounds=None,
features=None,
batch_shape=None,
dev_str=None):
"""
Create voxel grid :math:`\mathbf{X}_v\in\mathbb{R}^{x×y×z×(3+N+1)}` from homogeneous co-ordinates
:math:`\mathbf{X}_w\in\mathbb{R}^{num\_coords×4}`. Each voxel contains 3+N+1 values: the mean normalized
co-ordinate inside the voxel for the projected pixels with :math:`0 < x, y, z < 1`, N coordinte features (optional),
and also the number of projected pixels inside the voxel.
Grid resolutions and dimensions are also returned separately for each entry in the batch.
Note that the final batched voxel grid returned uses the maximum grid dimensions across the batch, this means
some returned grids may contain redundant space, with all but the single largest batched grid occupying a subset
of the grid space, originating from the corner of minimum :math:`x,y,z` values.\n
`[reference] <https://en.wikipedia.org/wiki/Voxel>`_
:param coords: Homogeneous co-ordinates *[batch_shape,c,4]*
:type coords: array
:param voxel_shape_spec: Either the number of voxels in x,y,z directions, or the resolutions (metres) in x,y,z
directions, depending on mode. Batched or unbatched. *[batch_shape,3]* or *[3]*
:type voxel_shape_spec: array
:param mode: Shape specification mode, either "DIMS" or "RES"
:type mode: str
:param coord_bounds: Co-ordinate x, y, z boundaries *[batch_shape,6]* or *[6]*
:type coord_bounds: array
:param features: Co-ordinate features *[batch_shape,c,4]*.
E.g. RGB values, low-dimensional features, etc.
Features mapping to the same voxel are averaged.
:type features: array
:param batch_shape: Shape of batch. Inferred from inputs if None.
:type batch_shape: sequence of ints, optional
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Same as x if None.
:type dev_str: str, optional
:return: Voxel grid *[batch_shape,x_max,v_max,z_max,3+feature_size+1]*, dimensions *[batch_shape,3]*, resolutions *[batch_shape,3]*, voxel_grid_lower_corners *[batch_shape,3]*
"""
if batch_shape is None:
batch_shape = coords.shape[:-2]
if dev_str is None:
dev_str = _ivy.dev_str(coords)
# shapes as list
batch_shape = list(batch_shape)
num_batch_dims = len(batch_shape)
num_coords_per_batch = coords.shape[-2]
# voxel shape spec as array
if len(voxel_shape_spec) is 3:
# BS x 1 x 3
voxel_shape_spec = _ivy.expand_dims(
_ivy.tile(
_ivy.reshape(_ivy.array(voxel_shape_spec),
[1] * num_batch_dims + [3]), batch_shape + [1]),
-2)
# coord bounds spec as array
if coord_bounds is not None and len(coord_bounds) is 6:
# BS x 1 x 6
coord_bounds = _ivy.expand_dims(
_ivy.tile(
_ivy.reshape(_ivy.array(coord_bounds, dtype_str='float32'),
[1] * num_batch_dims + [6]), batch_shape + [1]),
-2)
# BS x N x 3
coords = coords[..., 0:3]
if coord_bounds is not None:
# BS x 1
x_min = coord_bounds[..., 0:1]
y_min = coord_bounds[..., 1:2]
z_min = coord_bounds[..., 2:3]
x_max = coord_bounds[..., 3:4]
y_max = coord_bounds[..., 4:5]
z_max = coord_bounds[..., 5:6]
# BS x N x 1
x_coords = coords[..., 0:1]
y_coords = coords[..., 1:2]
z_coords = coords[..., 2:3]
x_validity_mask = _ivy.logical_and(x_coords > x_min, x_coords < x_max)
y_validity_mask = _ivy.logical_and(y_coords > y_min, y_coords < y_max)
z_validity_mask = _ivy.logical_and(z_coords > z_min, z_coords < z_max)
# BS x N
full_validity_mask = _ivy.logical_and(
_ivy.logical_and(x_validity_mask, y_validity_mask),
z_validity_mask)[..., 0]
# BS x 1 x 3
bb_mins = coord_bounds[..., 0:3]
bb_maxs = coord_bounds[..., 3:6]
bb_ranges = bb_maxs - bb_mins
else:
# BS x N
full_validity_mask = _ivy.cast(
_ivy.ones(batch_shape + [num_coords_per_batch]), 'bool')
# BS x 1 x 3
bb_mins = _ivy.reduce_min(coords, axis=-2, keepdims=True)
bb_maxs = _ivy.reduce_max(coords, axis=-2, keepdims=True)
bb_ranges = bb_maxs - bb_mins
# get voxel dimensions
if mode is 'DIMS':
# BS x 1 x 3
dims = _ivy.cast(voxel_shape_spec, 'int32')
elif mode is 'RES':
# BS x 1 x 3
res = _ivy.cast(voxel_shape_spec, 'float32')
dims = _ivy.cast(_ivy.ceil(bb_ranges / (res + MIN_DENOMINATOR)),
'int32')
else:
raise Exception(
'Invalid mode selection. Must be either "DIMS" or "RES"')
dims_m_one = _ivy.cast(dims - 1, 'int32')
# BS x 1 x 3
res = bb_ranges / (_ivy.cast(dims, 'float32') + MIN_DENOMINATOR)
# BS x NC x 3
voxel_indices = _ivy.minimum(
_ivy.cast(_ivy.floor((coords - bb_mins) / (res + MIN_DENOMINATOR)),
'int32'), dims_m_one)
# BS x NC x 3
voxel_values = ((coords - bb_mins) % res) / (res + MIN_DENOMINATOR)
feature_size = 0
if features is not None:
feature_size = features.shape[-1]
voxel_values = _ivy.concatenate([voxel_values, features], axis=-1)
# TNVC x len(BS)+1
valid_coord_indices = _ivy.cast(_ivy.indices_where(full_validity_mask),
'int32')
# scalar
total_num_valid_coords = valid_coord_indices.shape[0]
# TNVC x 3
voxel_values_pruned_flat = _ivy.gather_nd(voxel_values,
valid_coord_indices)
voxel_indices_pruned_flat = _ivy.gather_nd(voxel_indices,
valid_coord_indices)
# TNVC x len(BS)+2
if num_batch_dims == 0:
all_indices_pruned_flat = voxel_indices_pruned_flat
else:
batch_indices = valid_coord_indices[..., :-1]
all_indices_pruned_flat = _ivy.concatenate(
[batch_indices] + [voxel_indices_pruned_flat], -1)
# TNVC x 4
voxel_values_pruned_flat =\
_ivy.concatenate((voxel_values_pruned_flat, _ivy.ones([total_num_valid_coords, 1], dev_str=dev_str)), -1)
# get max dims list for scatter
if num_batch_dims > 0:
max_dims = _ivy.reduce_max(_ivy.reshape(dims, batch_shape + [3]),
axis=list(range(num_batch_dims)))
else:
max_dims = _ivy.reshape(dims, batch_shape + [3])
batch_shape_array_list = [_ivy.array(batch_shape, 'int32')
] if num_batch_dims != 0 else []
total_dims_list = _ivy.to_list(
_ivy.concatenate(
batch_shape_array_list +
[max_dims, _ivy.array([4 + feature_size], 'int32')], -1))
# BS x x_max x y_max x z_max x 4
scattered = _ivy.scatter_nd(
all_indices_pruned_flat,
voxel_values_pruned_flat,
total_dims_list,
reduction='replace' if _ivy.backend == 'mxnd' else 'sum')
# BS x x_max x y_max x z_max x 4 + feature_size, BS x 3, BS x 3, BS x 3
return _ivy.concatenate(
(scattered[..., :-1] /
(_ivy.maximum(scattered[..., -1:], 1.) + MIN_DENOMINATOR),
scattered[..., -1:]),
-1), dims[..., 0, :], res[..., 0, :], bb_mins[..., 0, :]
def ext_mat_and_intrinsics_to_cam_geometry_object(ext_mat, intrinsics, batch_shape=None, dev_str=None):
"""
Create camera geometry object from extrinsic matrix :math:`\mathbf{E}\in\mathbb{R}^{3×4}`, and camera intrinsics
object.
:param ext_mat: Extrinsic matrix *[batch_shape,3,4]*
:type ext_mat: array
:param intrinsics: camera intrinsics object
:type intrinsics: camera_intrinsics
:param batch_shape: Shape of batch. Inferred from inputs if None.
:type batch_shape: sequence of ints, optional
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Same as x if None.
:type dev_str: str, optional
:return: Camera geometry object
"""
if batch_shape is None:
batch_shape = ext_mat.shape[:-2]
if dev_str is None:
dev_str = _ivy.dev_str(ext_mat)
# shapes as list
batch_shape = list(batch_shape)
# num batch dims
num_batch_dims = len(batch_shape)
# BS x 4 x 4
ext_mat_homo = \
_ivy.concatenate(
(ext_mat, _ivy.tile(_ivy.reshape(_ivy.array([0., 0., 0., 1.], dev_str=dev_str),
[1] * (num_batch_dims + 1) + [4]),
batch_shape + [1, 1])), -2)
# BS x 4 x 4
inv_ext_mat_homo = _ivy.inv(ext_mat_homo)
# BS x 3 x 4
inv_ext_mat = inv_ext_mat_homo[..., 0:3, :]
# BS x 3 x 1
cam_center = inv_ext_mat_to_camera_center(inv_ext_mat)
# BS x 3 x 3
Rs = ext_mat[..., 0:3]
# BS x 3 x 3
inv_Rs = inv_ext_mat[..., 0:3]
# extrinsics object
extrinsics = _Extrinsics(cam_center, Rs, inv_Rs, ext_mat_homo, inv_ext_mat_homo)
# BS x 3 x 4
full_mat = calib_and_ext_to_full_mat(intrinsics.calib_mats, ext_mat)
# BS x 4 x 4
full_mat_homo = \
_ivy.concatenate((
full_mat, _ivy.tile(_ivy.reshape(_ivy.array([0., 0., 0., 1.], dev_str=dev_str),
[1] * (num_batch_dims + 1) + [4]),
batch_shape + [1, 1])), -2)
# BS x 4 x 4
inv_full_mat_homo = _ivy.inv(full_mat_homo)
# camera geometry object
cam_geometry = _CameraGeometry(intrinsics, extrinsics, full_mat_homo, inv_full_mat_homo)
return cam_geometry
def rasterize_triangles(pixel_coords_triangles, image_dims, batch_shape=None, dev_str=None):
"""
Rasterize image-projected triangles
based on: https://www.scratchapixel.com/lessons/3d-basic-rendering/rasterization-practical-implementation/rasterization-stage
and: https://www.scratchapixel.com/lessons/3d-basic-rendering/rasterization-practical-implementation/rasterization-practical-implementation
:param pixel_coords_triangles: Projected image-space triangles to be rasterized
*[batch_shape,input_size,3,3]*
:type pixel_coords_triangles: array
:param image_dims: Image dimensions.
:type image_dims: sequence of ints
:param batch_shape: Shape of batch. Inferred from Inputs if None.
:type batch_shape: sequence of ints, optional
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Same as x if None.
:type dev_str: str, optional
:return: Rasterized triangles
"""
if batch_shape is None:
batch_shape = []
if dev_str is None:
dev_str = _ivy.dev_str(pixel_coords_triangles)
# shapes as list
batch_shape = list(batch_shape)
num_batch_dims = len(batch_shape)
image_dims = list(image_dims)
input_image_dims = pixel_coords_triangles.shape[num_batch_dims:-2]
input_image_dims_prod = _reduce(_mul, input_image_dims, 1)
# BS x 3 x 2
pixel_xy_coords = pixel_coords_triangles[..., 0:2]
# BS x 3 x 1
pixel_x_coords = pixel_coords_triangles[..., 0:1]
pixel_y_coords = pixel_coords_triangles[..., 1:2]
# 1
x_min = _ivy.reshape(_ivy.reduce_min(pixel_x_coords, keepdims=True), (-1,))
x_max = _ivy.reshape(_ivy.reduce_max(pixel_x_coords, keepdims=True), (-1,))
x_range = x_max - x_min
y_min = _ivy.reshape(_ivy.reduce_min(pixel_y_coords, keepdims=True), (-1,))
y_max = _ivy.reshape(_ivy.reduce_max(pixel_y_coords, keepdims=True), (-1,))
y_range = y_max - y_min
# 2
bbox = _ivy.concatenate((x_range, y_range), 0)
img_bbox_list = [int(item) for item in _ivy.to_list(_ivy.concatenate((y_range + 1, x_range + 1), 0))]
# BS x 2
v0 = pixel_xy_coords[..., 0, :]
v1 = pixel_xy_coords[..., 1, :]
v2 = pixel_xy_coords[..., 2, :]
tri_centres = (v0 + v1 + v2) / 3
# BS x 1
v0x = v0[..., 0:1]
v0y = v0[..., 1:2]
v1x = v1[..., 0:1]
v1y = v1[..., 1:2]
v2x = v2[..., 0:1]
v2y = v2[..., 1:2]
# BS x BBX x BBY x 2
uniform_sample_coords = _ivy_svg.create_uniform_pixel_coords_image(img_bbox_list, batch_shape)[..., 0:2]
P = _ivy.round(uniform_sample_coords + tri_centres - bbox / 2)
# BS x BBX x BBY x 1
Px = P[..., 0:1]
Py = P[..., 1:2]
v0v1_edge_func = ((Px - v0x) * (v1y - v0y) - (Py - v0y) * (v1x - v0x)) >= 0
v1v2_edge_func = ((Px - v1x) * (v2y - v1y) - (Py - v1y) * (v2x - v1x)) >= 0
v2v0_edge_func = ((Px - v2x) * (v0y - v2y) - (Py - v2y) * (v0x - v2x)) >= 0
edge_func = _ivy.logical_and(_ivy.logical_and(v0v1_edge_func, v1v2_edge_func), v2v0_edge_func)
batch_indices_list = list()
for i, batch_dim in enumerate(batch_shape):
# get batch shape
batch_dims_before = batch_shape[:i]
num_batch_dims_before = len(batch_dims_before)
batch_dims_after = batch_shape[i + 1:]
num_batch_dims_after = len(batch_dims_after)
# [batch_dim]
batch_indices = _ivy.arange(batch_dim, dtype_str='int32', dev_str=dev_str)
# [1]*num_batch_dims_before x batch_dim x [1]*num_batch_dims_after x 1 x 1
reshaped_batch_indices = _ivy.reshape(batch_indices, [1] * num_batch_dims_before + [batch_dim] +
[1] * num_batch_dims_after + [1, 1])
# BS x N x 1
tiled_batch_indices = _ivy.tile(reshaped_batch_indices, batch_dims_before + [1] + batch_dims_after +
[input_image_dims_prod * 9, 1])
batch_indices_list.append(tiled_batch_indices)
# BS x N x (num_batch_dims + 2)
all_indices = _ivy.concatenate(
batch_indices_list + [_ivy.cast(_ivy.flip(_ivy.reshape(P, batch_shape + [-1, 2]), -1),
'int32')], -1)
# offset uniform images
return _ivy.cast(_ivy.flip(_ivy.scatter_nd(_ivy.reshape(all_indices, [-1, num_batch_dims + 2]),
_ivy.reshape(_ivy.cast(edge_func, 'int32'), (-1, 1)),
batch_shape + image_dims + [1],
reduction='replace' if _ivy.backend == 'mxnd' else 'sum'), -3), 'bool')
def compute_link_matrices(self, joint_angles, link_num, batch_shape=None):
"""
Compute homogeneous transformation matrices relative to base frame, up to link_num of links.
:param joint_angles: Joint angles of the robot *[batch_shape,num_joints]*
:type joint_angles: array
:param link_num: Link number for which to compute matrices up to
:type link_num: int
:param batch_shape: Shape of batch. Inferred from inputs if None.
:type batch_shape: sequence of ints, optional
:return: The link_num matrices, up the link_num *[batch_shape,link_num,4,4]*
"""
if batch_shape is None:
batch_shape = joint_angles.shape[:-1]
batch_shape = list(batch_shape)
num_batch_dims = len(batch_shape)
# BS x 1 x NJ
try:
dh_joint_angles = _ivy.expand_dims(joint_angles * self._dh_joint_scales - self._dh_joint_offsets, -2)
except:
d = 0
# BS x 1 x 4 x 4
A00 = _ivy.identity(4, batch_shape=batch_shape + [1])
Aitoip1dashs = list()
Aiip1s = list()
A0is = [A00]
# repeated blocks
# BS x 1 x NJ
dis = _ivy.tile(_ivy.reshape(self._dis, [1] * num_batch_dims + [1, self._num_joints]),
batch_shape + [1, 1])
# BS x 1 x 4
bottom_row = _ivy.tile(
_ivy.reshape(_ivy.array([0., 0., 0., 1.]), [1] * num_batch_dims + [1, 4]),
batch_shape + [1, 1])
# BS x 1 x 3
start_of_bottom_middle = _ivy.tile(
_ivy.reshape(_ivy.array([0., 0., 1.]), [1] * num_batch_dims + [1, 3]),
batch_shape + [1, 1])
# BS x 1 x 2
zeros = _ivy.zeros(batch_shape + [1, 2])
for i in range(self._num_joints):
# BS x 1 x 4
top_row = _ivy.concatenate((_ivy.cos(dh_joint_angles[..., i:i + 1]),
-_ivy.sin(dh_joint_angles[..., i:i + 1]), zeros), -1)
top_middle_row = _ivy.concatenate((_ivy.sin(dh_joint_angles[..., i:i + 1]),
_ivy.cos(dh_joint_angles[..., i:i + 1]), zeros), -1)
bottom_middle_row = _ivy.concatenate((start_of_bottom_middle, dis[..., i:i + 1]), -1)
# BS x 4 x 4
Aitoip1dash = _ivy.concatenate((top_row, top_middle_row, bottom_middle_row, bottom_row), -2)
# (BSx4) x 4
Aitoip1dash_flat = _ivy.reshape(Aitoip1dash, (-1, 4))
# (BSx4) x 4
Aiip1_flat = _ivy.matmul(Aitoip1dash_flat, self._AidashtoAis[i + 1])
# BS x 4 x 4
Aiip1 = _ivy.reshape(Aiip1_flat, batch_shape + [4, 4])
# BS x 4 x 4
A0ip1 = _ivy.matmul(A0is[-1][..., 0, :, :], Aiip1)
# append term to lists
Aitoip1dashs.append(Aitoip1dash)
Aiip1s.append(Aiip1)
A0is.append(_ivy.expand_dims(A0ip1, -3))
if i + 1 == link_num:
# BS x LN x 4 x 4
return _ivy.concatenate(A0is, -3)
raise Exception('wrong parameter entered for link_num, please enter integer from 1-' + str(self._num_joints))
