def test_subtensor(mode): items = [[0, True, True, True, False], [1, False, False, False, True]] data = [Tensor(np.random.random((5, 5))), Tensor(np.random.random(2))] if mode == "get": op = builtin.Subtensor(items) data = data[:1] if mode == "set": op = builtin.SetSubtensor(items) if mode == "inc": op = builtin.IncrSubtensor(items) tensors = [Tensor(0), Tensor(4), Tensor(2), Tensor(3)] @trace(symbolic=True, capture_as_const=True) def fwd(*tensors): return apply(op, *tensors)[0] result = fwd(*data, *tensors) check_pygraph_dump(fwd, data + tensors, [result])
def set_subtensor(input, value, tuple_val): input, tensors, items = unpack_getitem(input, tuple_val) op = builtin.SetSubtensor(items) return invoke_op(op, (input, value, *tensors))