def test_equals_argmax(self): """Validates that the =(argmax, argmax) composition works.""" lhs = tile.Value.from_ndims(2) rhs = tile.Value.from_ndims(2) out = op.equal(op.argmax(lhs), op.argmax(rhs)) func = tile.compose(self._ctx, self._dev, inputs=[('lhs', lhs), ('rhs', rhs)], outputs=[('out', out)]) invoker = plaidml.Invoker(self._ctx, func) invoker.set_input('lhs', self.make_inited_tensor((3, 4))) invoker.set_input('rhs', self.make_inited_tensor((3, 4))) output = self.make_output_tensor(invoker.get_output_shape('out')) invoker.set_output('out', output) invoker.invoke() with output.mmap_current() as view: for dim in range(0, 2): self.assertEqual(view[(dim, 0)], True)
def test_argmax(self): """Validates that the composition works.""" inp = tile.Value.from_ndims(2) out = op.argmax(inp) func = tile.compose(self._ctx, self._dev, inputs=[('inp', inp)], outputs=[('out', out)]) invoker = plaidml.Invoker(self._ctx, func) invoker.set_input('inp', self.make_inited_tensor((3, 4))) output = self.make_output_tensor(invoker.get_output_shape('out')) invoker.set_output('out', output) invoker.invoke() # Any increasing tensor reduced along an increasing dimension will # have an argmax equal to the dimension size with output.mmap_current() as view: for dim in range(0, 3): self.assertEqual(view[(dim)], 3)
def argmin(data, axis=-1, keepdims=1): if not keepdims: raise NotImplementedError( 'ArgMin with keepdims=0 is not yet implemented by the PlaidML ONNX backend' ) return (op.argmax(-data, axis=axis), )