jacinle.nd package

Submodules

jacinle.nd.batch module

jacinle.nd.batch.batchify(inputs: Sequence[Any]) → Any[source]

Recursively combine a list of inputs into a batch. This function handles tuples, lists, dicts, and numpy arrays.

Examples

>>> batchify([np.array([1, 2, 3]), np.array([4, 5, 6])])
array([[1, 2, 3],
        [4, 5, 6]])
>>> batchify([
...     {'a': np.array([1, 2, 3]), 'b': np.array([4, 5, 6])},
...     {'a': np.array([7, 8, 9]), 'b': np.array([10, 11, 12])},
... ])
{'a': array([[1, 2, 3],
            [7, 8, 9]]),
'b': array([[ 4,  5,  6],
            [10, 11, 12]])}
Parameters:inputs – a list of inputs to be batched.
Returns:a batched input.
jacinle.nd.batch.unbatchify(inputs: Any) → List[Any][source]

Recursively split a batch into a list of inputs. This function handles tuples, lists, dicts, and numpy arrays. This function is the inverse of batchify().

Example

>>> unbatchify(np.array([[1, 2, 3], [4, 5, 6]]))
[array([1, 2, 3]), array([4, 5, 6])]
>>> unbatchify({'a': np.array([[1, 2, 3], [4, 5, 6]]), 'b': np.array([[7, 8, 9], [10, 11, 12]])})
[{'a': array([1, 2, 3]), 'b': array([7, 8, 9])}, {'a': array([4, 5, 6]), 'b': array([10, 11, 12])}]
Parameters:inputs – a batched input.
Returns:a list of inputs.

jacinle.nd.indexing module

jacinle.nd.indexing.one_hot(label: numpy.ndarray, nr_classes: int, dtype='float32') → numpy.ndarray[source]

Convert a label array to one-hot array. This function works for either 0d (scalar) or 1d (vector) label array. If you want to convert higher dimensional label array, use one_hot_nd() instead.

Parameters:
  • label – the label array.
  • nr_classes – the number of classes.
  • dtype – the data type of the one-hot array.
Returns:

the one-hot array.

jacinle.nd.indexing.one_hot_nd(label, nr_classes, dtype='float32')[source]

Convert a label array to one-hot array.

Parameters:
  • label – the label array.
  • nr_classes – the number of classes.
  • dtype – the data type of the one-hot array.
Returns:

the one-hot array.

jacinle.nd.indexing.index_select_batch(data: Union[numpy.ndarray, Sequence[numpy.ndarray]], indices: Union[Sequence[int], numpy.ndarray]) → numpy.ndarray[source]

Gather indices as batch indices from data, which can either be typical nd array or a list of nd array.

Parameters:
  • data – the data array.
  • indices – the indices to be selected.
Returns:

the selected data.

jacinle.nd.meta module

jacinle.nd.meta.isndarray(thing: Any) → bool[source]

Check if the given object is a numpy array.

jacinle.nd.meta.is_ndarray(thing: Any) → bool[source]
jacinle.nd.meta.nd_concat(list_of_arrays: Sequence[numpy.ndarray]) → Optional[numpy.ndarray][source]

Concatenate a list of numpy arrays. This function handles the case when the list is empty or contains only one element.

Parameters:list_of_arrays – a list of numpy arrays.
Returns:the concatenated array, or None if the list is empty.
jacinle.nd.meta.nd_len(thing: Any) → int[source]

Get the length of a numpy array. This function handles the case when the input is a scalar or plain Python objects.

Parameters:thing – the input array.
Returns:the length of the array, or 1 if the input is a scalar or plain Python objects.
jacinle.nd.meta.nd_batch_size(thing: Any) → int[source]

Get the batch size of a numpy array. This function handles the case when the input a nested list or dict.

Examples

>>> nd_batch_size(np.array([1, 2, 3]))
3
>>> nd_batch_size([np.zeros((2, 3)), np.zeros((2, 5))])
2
>>> nd_batch_size({'a': np.zeros((2, 3)), 'b': np.zeros((2, 5))})
2
Parameters:thing – the input array or nested list/dict.
Returns:the batch size of the array.
jacinle.nd.meta.nd_split_n(ndarray: numpy.ndarray, n: int) → List[numpy.ndarray][source]

Split a numpy array into n parts. If the size is not divisible by n, the last part will be larger.

Parameters:
  • ndarray – the array to be split.
  • n – the number of parts.
Returns:

a list of arrays.

jacinle.nd.meta.size_split_n(full_size: Optional[int], n: int) → Optional[List[int]][source]

Split a size into n parts. If the size is not divisible by n, the last part will be larger. When the size is None, None will be returned.

Parameters:
  • full_size – the size to be split.
  • n – the number of parts.
Returns:

a list of sizes.

jacinle.nd.shape module

jacinle.nd.shape.unsqueeze(x: numpy.ndarray, *axes) → numpy.ndarray[source]

Unsqueeze the numpy array at the given axes. Similar to the PyTorch’s torch.unsqueeze().

Parameters:
  • x – the input array.
  • axes – the axes to unsqueeze.
Returns:

the unsqueezed array.

jacinle.nd.shape.unsqueeze_as(x: numpy.ndarray, y: numpy.ndarray, *x_axes_in_y) → numpy.ndarray[source]

Unsqueeze the numpy array x as the shape of y. The corresponding axes in x are specified by x_axes_in_y.

Example

>>> x = np.zeros((2, 3, 4))
>>> y = np.zeros((2, 3, 4, 5))
>>> unsqueeze_as(x, y, 0, 1, 2).shape
(2, 3, 4, 1)
Parameters:
  • x – the input array.
  • y – the target array.
  • x_axes_in_y – the axes in x that correspond to the axes in y.
Returns:

the unsqueezed array.

jacinle.nd.shape.softmax(x: numpy.ndarray, axis: int = -1) → numpy.ndarray[source]

Compute the softmax of the input array at the given axis.

\[\text{softmax}(x)_i = \frac{e^{x_i}}{\sum_j e^{x_j}}\]
Parameters:
  • x – the input array.
  • axis – the axis to apply softmax.
Returns:

the softmax array.