def softmax(t: Tensor, axis: int) -> Tensor:
"""
Computes the softmax operation on a Tensor. This recales the slices of axis
such that all elements are within [0, 1] and sum to 1.
The output shape and dtype matches the input.
Args:
t: Tensor
Tensor to be softmaxed.
axis: int
The axis along which the softmax will be computed.
Returns:
out: Tensor
The softmaxed tensor
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
settings = ctx._get_op_settings('softmax')
opid = _ir.OperatorIdentifier("ai.onnx", "Softmax", 11, _ir.NumInputs(
1, 1), 1)
op = pb_g.createConnectedOp_SoftmaxOp(
{0: t.id}, {0: g._create_tensor_id(f"softmax_out")}, opid,
handle_negative_axis(t, axis), settings)
return Tensor._from_pb_tensor(op.outTensor(0))
def mul(lhs: Tensor, rhs: Tensor) -> Tensor:
"""
Multiplies two Tensors element-wise.
Follows numpy broadcasting rules.
Arguments must have the same dtype.
Args:
lhs, rhs: Tensor
Tensors to be multiplied.
Returns:
mul: Tensor
The product of lhs and rhs
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, lhs, rhs)
settings = ctx._get_op_settings('mul')
opid = _ir.OperatorIdentifier("ai.onnx", "Mul", 7, _ir.NumInputs(2, 2), 1)
op = pb_g.createConnectedOp_MulOp(
{
0: lhs.id,
1: rhs.id
},
{
0: g._create_tensor_id("mul_out"),
},
opid,
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def cast(t: Tensor, data_type: dtype) -> Tensor:
"""
Casts tensor `t` to data type `dtype`.
Args:
t: Tensor
Tensors to be casted.
data_type: popart.ir.dtypes.dtype
Dtype to cast to
Returns:
add: Tensor
The sum of lhs and rhs
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
settings = ctx._get_op_settings('cast')
opid = _ir.OperatorIdentifier("ai.onnx", "Cast", 9, _ir.NumInputs(1, 1), 1)
op = pb_g.createConnectedOp_CastOp(
{0: t.id},
{0: g._create_tensor_id(f"{t.id}_{data_type._name}")},
_to=data_type._pb_dtype,
opid=opid,
settings=settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def gelu(t: Tensor) -> Tensor:
"""
Computes the Gelu activation on a Tensor.
https://arxiv.org/abs/1606.08415
Args:
t: Tensor
Input tensor.
Returns:
out: Tensor
Output tensor.
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
settings = ctx._get_op_settings('gelu')
opid = _ir.OperatorIdentifier("ai.graphcore", "Gelu", 1,
_ir.NumInputs(1, 1), 1)
op = pb_g.createConnectedOp_GeluOp({0: t.id},
{0: g._create_tensor_id(f"gelu_out")},
opid, settings)
return Tensor._from_pb_tensor(op.outTensor(0))
def get_expected_forward_inputs_from_call(
call_info: CallInfo,
grad_info: GradGraphInfo) -> Mapping[Tensor, Tensor]:
"""Utility function for constructing inputs to calling a grad graph.
Args:
call_info: `popart.ir.ops.call.CallInfo`
Callsite info of a call to the graph that was auto-differentiated. This can be accessed by
using `ops.call_with_info()`
grad_info: `GradGraphInfo`
Output of autodiff on a graph.
Returns: `Mapping[Tensor, Tensor]`
from: a Tensor in the gradient Graph
to: an input or output tensor at a callsite of the corresponding forward Graph.
"""
if call_info.called_graph != grad_info.forward_graph:
raise TypeError(
"The called graph does not match the graph that was auto-differentiated."
)
return {
Tensor._from_pb_tensor(grad_info.graph._pb_graph.getInputTensor(idx)):
call_info.subgraph_to_op_tensor(act.fwd_tensor)
for idx, act in enumerate(grad_info.expected_inputs)
if act.connection_type == ExpectedConnectionType.Fwd
}
def transpose(t: Tensor,
permutation: Optional[Tuple[int, ...]] = None) -> Tensor:
"""
Permute the axes of a Tensor. By default reverses the axes of t.
Args:
t: Tensor
Tensor to be transposed.
permutation: tuple of ints (optional)
Tuple containing the a permutation of [0, N-1] where N is the
rank of input `t`. If not provided, the axes will be reversed.
Returns:
out: Tensor
The transposed tensor
"""
permutation = _handle_permuation(t, permutation)
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
settings = ctx._get_op_settings('transpose')
opid = _ir.OperatorIdentifier("ai.onnx", "Transpose", 1,
_ir.NumInputs(1, 1), 1)
op = pb_g.createConnectedOp_TransposeOp(
{0: t.id},
{0: g._create_tensor_id(f"{t.name}_T")},
opid,
permutation,
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def equal(lhs: Tensor, rhs: Tensor) -> Tensor:
"""
Compares two Tensors element-wise with an equal operator.
Follows numpy broadcasting rules.
Args:
lhs, rhs: Tensor
Tensors to be compared.
Returns:
out: Tensor
The value (lhs == rhs)
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, lhs, rhs)
settings = ctx._get_op_settings('equal')
opid = _ir.OperatorIdentifier("ai.onnx", "Equal", 7, _ir.NumInputs(2, 2),
1)
op = pb_g.createConnectedOp_AndOp(
{
0: lhs.id,
1: rhs.id
},
{
0: g._create_tensor_id("equal_out"),
},
opid,
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def increment_mod_(t: Tensor, increment: float, modulus: float) -> Tensor:
"""
Compute `(t + increment) % modulus` inplace on `t`
Args:
t: Tensor
Tensor to increment (modulo)
increment: float
How much to increment the input tensor by.
increment: float
The modulo operand.
Returns:
out: Tensor (alias of input `t`)
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
settings = ctx._get_op_settings('increment_mod_inplace')
op = pb_g.createConnectedOp_IncrementModInplaceOp(
{
0: t.id,
},
{
0: g._create_tensor_id("increment_mod_inplace_out"),
},
increment,
modulus,
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def logical_not(t: Tensor) -> Tensor:
"""
Computes element-wise the value of NOT t.
Inputs will be cast to bool if needed.
Args:
t: Tensor
Input tensor.
Returns:
out: Tensor
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
t = cast_if_needed(t, dtypes.bool)
settings = ctx._get_op_settings('not')
opid = _ir.OperatorIdentifier("ai.onnx", "Not", 1, _ir.NumInputs(1, 1), 1)
op = pb_g.createConnectedOp_NotOp(
{0: t.id},
{0: g._create_tensor_id("not_out")},
opid,
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def replicated_all_gather(t: Tensor,
group: Optional[CommGroup] = None) -> Tensor:
"""Gathers tensor `t` across replicas. Output tensor contains in the values of `t` from each replica.
Args:
t (Tensor): Tensor to be reduced. Must be rank=1.
group (Optional[CommGroup], optional): Replicas to gather from. Defaults to All replicas.
Returns:
Tensor: Gathered tensor.
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
if group is None:
group = CommGroup()
settings = ctx._get_op_settings('replicated_all_gathered')
opid = _ir.OperatorIdentifier("ai.graphcore", "ReplicatedAllGather", 1,
_ir.NumInputs(1, 1), 1)
op = pb_g.createConnectedOp_ReplicatedAllGatherOp(
{0: t.id}, {0: g._create_tensor_id(t.name + "_all_gathered")}, opid,
group, settings)
return Tensor._from_pb_tensor(op.outTensor(0))
def replicated_reduce_scatter(t: Tensor,
op: CollectiveOperator = CollectiveOperator.Add,
group: Optional[CommGroup] = None) -> Tensor:
"""Reduces tensor `t` across replicas. Each replica will only receive a unique slice of `t`.
Args:
t (Tensor): Tensor to be reduced. Inputs will be flattened.
op (CollectiveOperator, optional): Operation to reduce with. Defaults to CollectiveOperator.Add.
group (Optional[CommGroup], optional): Replicas to reduce across. Defaults to All replicas.
Returns:
Tensor: A slice of the reduced tensor. Always a 1D tensor.
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
if group is None:
group = CommGroup()
settings = ctx._get_op_settings('replicated_reduce_scatter')
opid = _ir.OperatorIdentifier("ai.graphcore", "ReplicatedReduceScatter", 1,
_ir.NumInputs(1, 1), 1)
op = pb_g.createConnectedOp_ReplicatedReduceScatterOp(
{0: t.id}, {0: g._create_tensor_id(t.name + "_reduce_scattered")},
opid, op, group, settings)
return Tensor._from_pb_tensor(op.outTensor(0))
def sub(lhs: Tensor, rhs: Tensor) -> Tensor:
"""Subtracts two Tensors element-wise. Follows numpy broadcasting rules. Arguments must have the same dtype.
Args:
lhs, rhs: Tensor
Tensors to be subtracted.
Returns:
add: Tensor
The value of (lhs - rhs)"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, lhs, rhs)
settings = ctx._get_op_settings('sub')
opid = _ir.OperatorIdentifier("ai.onnx", "Sub", 7, _ir.NumInputs(2, 2), 1)
op = pb_g.createConnectedOp_SubtractOp(
{
0: lhs.id,
1: rhs.id
},
{
0: g._create_tensor_id("sub_out"),
},
opid,
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def detach_(x: Tensor) -> Tensor:
"""
This is the inplace version of :func:`~ops.detach`. Behaviour is the same, but blocks gradient
propagation inplace on the input tensor.
Args:
x: Tensor
Input tensor.
Returns:
out: Tensor
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, x)
settings = ctx._get_op_settings('detach_inplace')
op = pb_g.createConnectedOp_DetachInplaceOp(
{0: x.id},
{0: g._create_tensor_id("detach_out")},
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def detach(t: Tensor) -> Tensor:
"""
Prevents gradient computation of this tensor.
Numerically equivlent to the identity op.
Args:
t: Tensor
Input tensor.
Returns:
out: Tensor
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
settings = ctx._get_op_settings('detach')
opid = _ir.OperatorIdentifier("ai.graphcore", "Detach", 1,
_ir.NumInputs(1, 1), 1)
op = pb_g.createConnectedOp_DetachOp(
{0: t.id},
{0: g._create_tensor_id("detach_out")},
opid,
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def copy_var_update_(t: Tensor, X: Tensor) -> Tensor:
"""
Updates tensor `t` inplace by copying `X`.
Args:
t: Tensor
Tensor to be updated.
X: Tensor
Value to update the variable
Returns:
updated: Tensor
An alias to the variable.
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t, X)
settings = ctx._get_op_settings('copy_var_update')
op = pb_g.createConnectedOp_CopyVarUpdateOp(
{
0: t.id,
1: X.id
},
{0: g._create_tensor_id('updated__' + t.name)},
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def slice(t: Tensor,
start: Optional[Union[int, List[Optional[int]]]] = None,
stop: Optional[Union[int, List[Optional[int]]]] = None,
step: Optional[Union[int, List[Optional[int]]]] = None,
axis: Optional[Union[int, List[int]]] = None) -> Tensor:
"""
Selects elements from a tensor using a slice or multiple slices.
A slice specifies the start (inclusive) and stop (exclusive) index of elements to select.
Multiple slices can be specified using a list of items for each parameter (start, stop, step).
If step is `-1` the slice is performed backwards.
If axis is not specified, each slice will correspond to axis 0 to N where N is the number of slices.
Examples:
```
t == slice(t) == slice(t, axis=1)
slice(t, start=1) # Slice axis 0 from start index 1
slice(t, start=[1,2]) == slice(t, start=[1,2], axis=[0,1])
slice(t, stop=-2) # Slice axis 0 upto second last element (exclusive)
slice(t, stop=3, step=-1) # Slice backwards from last element (inclusive) to third last element (exclusive)
```
Args:
t (Tensor): Tensor to slice
start: Index of first element (inclusive) or `None` which defaults to 0.
stop: Index of last element (exclusive) or `None` which defaults to last
element (inclusive) if step is forward or first element (inclusive) if step is backwards.
step: `1` for forward or `-1` for backwards.
axis: Axis of tensor to slice on or `None` will default to each axis sequentially.
Returns:
Tensor: output tensor
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
if start is None and stop is None and step is None:
return t
start, stop, step, axis = process_args(start, stop, step, axis)
opid = _ir.OperatorIdentifier("ai.onnx", "Slice", 11, _ir.NumInputs(1, 1),
1)
settings = ctx._get_op_settings("slice")
op = pb_g.createConnectedOp_SliceOp(
{0: t.id},
{0: g._create_tensor_id("slice_out")},
starts_=start,
ends_=stop,
axes_=axis,
steps_=step,
opid=opid,
settings=settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def nll_loss_with_softmax_grad(
probs: Tensor,
labels: Tensor,
loss_grad: Union[float, Tensor] = 1,
reduction: popart.ReductionType = popart.ReductionType.Mean,
ignore_index: Optional[int] = None) -> Tuple[Tensor, Tensor]:
"""Computes the negative log likelihood loss `l` and returns the gradient `dE/dx` where `probs = softmax(x)`.
Argument `loss_grad` should be the gradient `dE/dl`, where `E` is the error from which back propagation is initialised from.
Typically `E = l`, so to return `dl/dx` then `loss_grad` should be `dl/dl` which would be `1`.
Args:
probs (Tensor): probabilities. Expected to be the ouput of `ops.softmax`.
labels (Tensor): labels. Target values for the probabilities.
loss_grad (Tensor): `dE/dl`. Supports float32 dtypes with float16 `probs`
reduction (ReductionType): Specify how to reduce the loss. Defaults to ReductionType.Mean.
ignore_index (Optional[int]): Specify label values that should not contribute to `l` or `dE/dx`. Defaults to None.
Returns:
Tuple[Tensor, Tensor]: (`l`, `dE/dx`)
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, probs, labels)
if not isinstance(loss_grad, Tensor):
loss_grad = constant(loss_grad, float32)
settings = ctx._get_op_settings('nll_loss_with_softmax_grad')
op = pb_g.createConnectedOp_NlllWithSoftmaxGradDirectOp(
{
0: probs.id,
1: labels.id,
2: loss_grad.id
},
{
0: g._create_tensor_id("loss"),
1: g._create_tensor_id("dx"),
},
convert_optional_int(ignore_index),
reduction,
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0)), Tensor._from_pb_tensor(
op.outTensor(1))
def dynamic_slice(t: Tensor, index: Tensor, axes: List[int], sizes: List[int],
no_overlap: bool) -> Tensor:
"""
Returns a cloned slice of the input Tensor.
The word "dynamic" refers to the fact that the index can be specified
during runtime.
A slice along an axis can be defined as by the tuple
( start, stop, step )
start - will be equal the index for the respective axis
stop - will be equal index + size for the respective axis
step - will equal 1
Limitations:
Assuming we would like to slice A with dimension (4, 3)
- Step other than 1 is not supported (i.e. t[::2,:] is not supported)
- Negative slicing is not supported (i.e. t[:-1,:] is not supported)
- stop greater than the size of the axis is not supported
(i.e. t[:5,:] is not supported)
Args:
t: Tensor
Input tensor.
index: Tensor
The indices to start the slice from.
axes: List[int]
The axes to slice from.
sizes: List[int]
The sizes of the slices for the specified axes.
For example:
If index = [1, 2], axes = [0, 3] and sizes = [2, 4], the Tensor will be sliced
t[1:2, :, :, 2:4]
no_overlap : bool
If set to true, then correct gradient backpropagation is only guaranteed if
each region in the output tensor has exactly one populator
(operation that writes data to this region).
There are no run-time or compile-time checks possible to ensure this.
Returns:
out: Tensor
A clone (i.e. not a view) of the sliced input tensor.
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t, index)
settings = ctx._get_op_settings('dynamicslice')
opid = _ir.OperatorIdentifier("ai.graphcore", "DynamicSlice", 1,
_ir.NumInputs(2, 2), 1)
op = pb_g.createConnectedOp_DynamicSliceOp(
{
0: t.id,
1: index.id
}, {0: g._create_tensor_id(f"dynamic_slice_out")}, opid, axes, sizes,
no_overlap, settings)
return Tensor._from_pb_tensor(op.outTensor(0))
def get_output_tensors(self) -> Tuple[Tensor, ...]:
"""Return outputs to the CallOp in index order.
Returns:
Tuple[Tensor, ...]
"""
return tuple(
Tensor._from_pb_tensor(t) for t in self._op.getOutputTensors())
def accumulate_mean_(t: Tensor, X: Tensor, step: Union[float,
Tensor]) -> Tensor:
"""
Updates a tensor `t` inplace using `t = (step/(step+1)) * t + (1/(step+1)) * X`.
Intended to be used to keep track of the mean of a series of values.
For example:
```
with g:
accum = pir.variable(0, dtype=pir.float32)
a = pir.variable(1, dtype=pir.float32)
b = pir.variable(2, dtype=pir.float32)
accumulate_mean(accum, a, 0.0)
accumulate_mean(accum, b, 1.0)
```
will result with `accum` having the value `(a+b)/2 = 1.5`.
Does not apply numpy broadcasting.
Uses mixed precision poplibs operations.
`t` and `X` must be the same shape, but can be different types.
`step` must be scalar.
Args:
`t`: Tensor
Tensor to be updated.
`X`: Tensor
Value to update the variable
step: Union[float, Tensor]]
Value representing the number of previously accumulated values.
Returns:
updated: Tensor
An alias to the variable.
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
step = t._ensure_tensor(step, dtype=dtypes.float32)
check_in_graph(g, t, X, step)
settings = ctx._get_op_settings('accumulate')
op = pb_g.createConnectedOp_AccumulateOp(
{
0: t.id,
1: X.id,
2: step.id
},
{
0: g._create_tensor_id('accumulate_mean__' + t.name),
},
_ir.AccumulationType.Mean,
_ir.OptimizerValue(),
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def tied_gather(t: Tensor,
indices: Tensor,
axis: int = 0,
available_memory_proportion: Optional[float] = None) -> Tensor:
"""
Select multiple elements from an array, given by `indices`, along a specified axis.
When `axis == 0`, it is equivlent to numpy "fancy indexing".
Numerically the same as the `gather` op but does not specify the tile
layout of the `indices` tensor. When preceding a `matmul` op the tile
layout of the indices is determined by the `matmul`, not the `tied_gather`.
This has a has lower memory footprint but costs extra cycles due to the exchange.
Pseudo example:
```
tied_gather(x, [1, 2, 3]) == [x[3], x[7], x[2]]
```
Args:
t: Tensor
Input tensor
indices: Tensor
The indices of the elements to extract
axis: int
Which axis to gather on. Default is 0.
available_memory_proportion: Optional[float]
The maximum proportion of available memory on each tile that this layer
should consume temporarily during the course of the operation.
Defaults to 1.0 if not set globally.
Returns:
gather: Tensor
The gathered elements concatenated.
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
check_in_graph(g, indices)
available_memory_proportion = convert_optional_float(
available_memory_proportion)
settings = ctx._get_op_settings("tiedgather")
op = pb_g.createConnectedOp_TiedGatherOp(
{
0: t.id,
1: indices.id
}, {0: g._create_tensor_id("tiedgather_out")},
axis_=axis,
available_memory_proportion_=available_memory_proportion,
settings=settings)
return Tensor._from_pb_tensor(op.outTensor(0))
def split(t: Tensor,
splits: Union[int, List[int]],
axis: int = 0) -> List[Tensor]:
"""
Splits a tensor on a given axis into a list of tensors.
Args:
t: Tensor
Tensor to be split.
splits: int or List[int]
Either an int which specifies the number of splits or a list of ints specifing the length of each output tensor.
axis: int (default 0)
Which axis to split on
Returns:
out: List[Tensor]
A list of tensors
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
axis = handle_negative_axis(t, axis)
if isinstance(splits, int):
axis_len = t.shape[axis]
if axis_len % splits != 0:
raise ValueError(
f"Split {splits} does not equally divide tensor axis {axis} of length {axis_len}."
)
splits = [axis_len // splits] * splits
outputs_t = {
i: g._create_tensor_id(f"{t.name}_split_{i}")
for i in range(len(splits))
}
settings = ctx._get_op_settings('split')
opid = _ir.OperatorIdentifier("ai.onnx", "Split", 2, _ir.NumInputs(1, 1),
1)
op = pb_g.createConnectedOp_SplitOp(
{0: t.id},
outputs_t,
axis_=axis,
split_=splits,
opid=opid,
settings=settings,
)
output = [
Tensor._from_pb_tensor(op.outTensor(i)) for i in range(len(splits))
]
return output
def get_input_tensors(self) -> Tuple[Tensor, ...]:
"""Get the input tensors to the graph.
Returns:
Tuple[Tensor, ...]: A tuple of all the input tensors.
"""
_pb_ins = self._pb_graph.getInputIds()
return tuple(
Tensor._from_pb_tensor(self._pb_graph.getTensor(o))
for o in _pb_ins)
def get_output_tensors(self) -> Tuple[Tensor, ...]:
"""Get the output tensors from the graph.
Returns:
Tuple[Tensor, ...]: A tuple of all the output tensors.
"""
_pb_outs = self._pb_graph.getOutputIds()
return tuple(
Tensor._from_pb_tensor(self._pb_graph.getTensor(o))
for o in _pb_outs)
def remote_load(
t: Tensor,
offset: Optional[Tensor] = None,
remote_buffer_handle: Optional[RemoteBufferHandle] = None) -> Tensor:
"""Load a tensor from remote (off-chip) buffer.
The tensor will be loaded from the memory location corresponding to
`remote_buffer_id` (specified in the `remote_buffer_handle`),
and will be stored in the memory location corresponding to `t`.
The relationship between `offset` and `remote_buffer_id` is thoroughly
described in `remote_store`.
See also: `remote_buffer_handle`, `remote_store`, `remote_load_`
Args:
t (Tensor): This tensor will be cloned, and the loaded data will written to the clone.
offset (Optional[Tensor], optional): Optional 0-rank Tensor.
Specify the row in the remote buffer the inTensor will be loaded from.
Defaults to None.
remote_buffer_handle (Optional[RemoteBufferHandle], optional): The handle to the remote
buffer. Defaults to None.
Returns:
Tensor: The tensor loaded from the remote buffer
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
if offset is not None:
check_in_graph(g, offset)
remote_buffer_handle = prepare_remote_buffer(t, remote_buffer_handle, g)
settings = ctx._get_op_settings('remote_load')
opid = _ir.OperatorIdentifier("ai.graphcore", "RemoteLoad", 1,
_ir.NumInputs(1, 2), 1)
if offset is not None:
op = pb_g.createConnectedOp_RemoteLoadOp(
{
0: t.id,
1: offset.id
}, {0: g._create_tensor_id("remote_load_out")}, opid, settings,
remote_buffer_handle.remote_buffer_id)
else:
op = pb_g.createConnectedOp_RemoteLoadOp(
{
0: t.id,
}, {0: g._create_tensor_id("remote_load_out")}, opid, settings,
remote_buffer_handle.remote_buffer_id)
return Tensor._from_pb_tensor(op.outTensor(0))
def gather(t: Tensor,
indices: Tensor,
axis: int = 0,
available_memory_proportion: Optional[float] = None) -> Tensor:
"""
Select multiple elements from an array, given by `indices`, along a specified axis.
When `axis == 0`, it is equivlent to numpy "fancy indexing".
Pseudo example:
```
gather(x, [1, 2, 3]) == [x[3], x[7], x[2]]
```
Args:
t: Tensor
Input tensor
indices: Tensor
The indices of the elements to extract
axis: int
Which axis to gather on. Default is 0.
available_memory_proportion: Optional[float]
The maximum proportion of available memory on each tile that this layer
should consume temporarily during the course of the operation.
Defaults to 1.0 if not set globally.
Returns:
gather: Tensor
The gathered elements concatenated.
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
check_in_graph(g, indices)
available_memory_proportion = convert_optional_float(
available_memory_proportion)
opid = _ir.OperatorIdentifier("ai.onnx", "Gather", 11, _ir.NumInputs(2, 2),
1)
settings = ctx._get_op_settings("gather")
op = pb_g.createConnectedOp_GatherOp(
{
0: t.id,
1: indices.id
}, {0: g._create_tensor_id("gather_out")},
opid=opid,
axis_=axis,
available_memory_proportion_=available_memory_proportion,
settings=settings)
return Tensor._from_pb_tensor(op.outTensor(0))
def remote_load_(
t: Tensor,
offset: Optional[Tensor] = None,
remote_buffer_handle: Optional[RemoteBufferHandle] = None) -> Tensor:
"""Load a tensor from remote (off-chip) buffer inplace.
This op is identical to `remote_load` with the exception of how `t` is handled.
In `remote_load` `t` is cloned and the output is written to the clone, whereas
in this version `t` is written to directly.
See also: `remote_buffer_handle`, `remote_store`, `remote_load`
Args:
t (Tensor): The tensor the loaded data will written to the clone.
offset (Optional[Tensor], optional): Optional 0-rank Tensor.
Specify the row in the remote buffer the inTensor will be loaded from.
Defaults to None.
remote_buffer_handle (Optional[RemoteBufferHandle], optional): The handle to the remote
buffer. Defaults to None.
Returns:
Tensor: The tensor loaded from the remote buffer
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
if offset is not None:
check_in_graph(g, offset)
remote_buffer_handle = prepare_remote_buffer(t, remote_buffer_handle, g)
settings = ctx._get_op_settings('remote_load_inplace')
opid = _ir.OperatorIdentifier("ai.graphcore", "RemoteLoadInplace", 1,
_ir.NumInputs(1, 2), 1)
if offset is not None:
op = pb_g.createConnectedOp_RemoteLoadInplaceOp(
{
0: t.id,
1: offset.id
}, {0: g._create_tensor_id("remote_load_inplace_out")}, opid,
settings, remote_buffer_handle.remote_buffer_id)
else:
op = pb_g.createConnectedOp_RemoteLoadInplaceOp(
{
0: t.id,
}, {0: g._create_tensor_id("remote_load_inplace_out")}, opid,
settings, remote_buffer_handle.remote_buffer_id)
return Tensor._from_pb_tensor(op.outTensor(0))
def ipu_copy(t: Tensor, destination: int,
source: Optional[int] = None) -> Tensor:
"""
Copies a Tensor to a virtual graph.
Args:
t: Tensor
Tensor to be copied.
destination: int
Ipu for the tensor to be copied to.
source: Optional[int]
Ipu for the tensor to be copied from.
By default, the source will be taken from the producer of `t`.
If `t` does not have a producer a source MUST be provided.
Returns:
t_copied: Tensor
The copied tensor
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
if source is None:
# Use internal method to infer the input tensor's virtual graph.
source = t._pb_tensor.getVirtualGraphIdUnsafe()
if source == -1:
raise ValueError(
f"Could not infer virtual graph for Tensor to be copied \"{t}\" . "
"Please specify `source` when copying for this tensor.")
settings = ctx._get_op_settings('ipucopy')
opid = _ir.OperatorIdentifier("ai.graphcore", "IpuCopy", 1,
_ir.NumInputs(1, 1), 1)
op = pb_g.createConnectedOp_IpuCopyOp(
{
0: t.id,
},
{
0: g._create_tensor_id(t.name + f"_c{destination}"),
},
opid,
source,
destination,
settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def slice_(t: Tensor,
start: Optional[Union[int, List[Optional[int]]]] = None,
stop: Optional[Union[int, List[Optional[int]]]] = None,
step: Optional[Union[int, List[Optional[int]]]] = None,
axis: Optional[Union[int, List[int]]] = None) -> Tensor:
"""
Selects elements from a tensor using a slice or multiple slices. Inplace.
This is the inplace version of :func:`~ops.slice`. Behaviour is the same, but modifies the
tensor inplace.
```
Args:
t (Tensor): Tensor to slice
start: Index of first element (inclusive) or `None` which defaults to 0.
stop: Index of last element (exclusive) or `None` which defaults to last
element (inclusive) if step is forward or first element (inclusive) if step is backwards.
step: `1` for forward or `-1` for backwards.
axis: Axis of tensor to slice on or `None` will default to each axis sequentially.
Returns:
Tensor: alias of the input tensor t.
"""
ctx = get_current_context()
g = ctx.graph
pb_g = g._pb_graph
check_in_graph(g, t)
if start is None and stop is None and step is None:
return t
start, stop, step, axis = process_args(start, stop, step, axis)
opid = _ir.OperatorIdentifier("ai.graphcore", "SliceInplace", 1,
_ir.NumInputs(1, 1), 1)
settings = ctx._get_op_settings("slice_inplace")
op = pb_g.createConnectedOp_SliceInplaceOp(
{0: t.id},
{0: g._create_tensor_id("slice_out")},
starts_=start,
ends_=stop,
axes_=axis,
steps_=step,
opid=opid,
settings=settings,
)
return Tensor._from_pb_tensor(op.outTensor(0))
def subgraph_to_op_tensor(self, subgraph_tensor: Tensor) -> Tensor:
"""
Provided an input or output tensor in the `called_graph`, this method
returns the associated input or output tensor on the CallOp.
Args:
subgraph_tensor (Tensor): The tensor in the subgraph.
Raises:
ValueError: If `subgraph_tensor` is not an input or output of the called graph.
Returns:
Tensor: The associated input or output tensor on the CallOp
"""
sgraph = self.called_graph._pb_graph
if sgraph.hasInputId(subgraph_tensor.id):
idx = sgraph.getInputIndex(subgraph_tensor.id)
return Tensor._from_pb_tensor(self._op.inTensor(idx))
if sgraph.hasOutputId(subgraph_tensor.id):
idx = sgraph.getOutputIndex(subgraph_tensor.id)
return Tensor._from_pb_tensor(self._op.outTensor(idx))
raise ValueError(
f"Tensor {subgraph_tensor.name} is not an Input or Output of the called graph {sgraph.id}"
)
