construct arguments to `Vec` according to axis index.
(value, axis_idx, x_value=0.0, dtype=np.float64, n_dims=3)
| 131 | |
| 132 | |
| 133 | def get_axis_params(value, axis_idx, x_value=0.0, dtype=np.float64, n_dims=3): |
| 134 | """construct arguments to `Vec` according to axis index.""" |
| 135 | zs = np.zeros((n_dims,)) |
| 136 | assert axis_idx < n_dims, "the axis dim should be within the vector dimensions" |
| 137 | zs[axis_idx] = 1.0 |
| 138 | params = np.where(zs == 1.0, value, zs) |
| 139 | params[0] = x_value |
| 140 | return list(params.astype(dtype)) |
| 141 | |
| 142 | |
| 143 | @torch.jit.script |