Method squeeze

ot/backend.py:1799–1800

(self, a, axis=None)

        ), f"Device discrepancy. First input is on {str(a_device)}, whereas second input is on {str(b_device)}"

    def squeeze(self, a, axis=None):
        return jnp.squeeze(a, axis=axis)

    def unsqueeze(self, a, axis):
        return jnp.expand_dims(a, axis=axis)