torchax/device-test-env.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py:354: UserWarning: Device capability of jax unspecified, assuming `cpu` and `cuda`. Please specify it via the `devices` argument of `register_backend`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.4167, -0.2754,  0.1437],\n",
       "        [-0.5498, -1.5257, -0.0837]])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Here's a CPU tensor\n",
    "a = torch.randn(2, 3, device='cpu')\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "NotImplementedError",
     "evalue": "Could not run 'aten::empty.memory_format' with arguments from the 'jax' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'aten::empty.memory_format' is only available for these backends: [CPU, Meta, QuantizedCPU, QuantizedMeta, MkldnnCPU, SparseCPU, SparseMeta, SparseCsrCPU, SparseCsrMeta, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradHIP, AutogradXLA, AutogradMPS, AutogradIPU, AutogradXPU, AutogradHPU, AutogradVE, AutogradLazy, AutogradMTIA, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, AutogradMeta, AutogradNestedTensor, Tracer, AutocastCPU, AutocastXPU, AutocastMPS, AutocastCUDA, FuncTorchBatched, BatchedNestedTensor, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PreDispatch, PythonDispatcher].\n\nCPU: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterCPU_1.cpp:2513 [kernel]\nMeta: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterMeta_0.cpp:5433 [kernel]\nQuantizedCPU: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterQuantizedCPU_0.cpp:311 [kernel]\nQuantizedMeta: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterQuantizedMeta_0.cpp:116 [kernel]\nMkldnnCPU: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterMkldnnCPU_0.cpp:230 [kernel]\nSparseCPU: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterSparseCPU_0.cpp:843 [kernel]\nSparseMeta: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterSparseMeta_0.cpp:187 [kernel]\nSparseCsrCPU: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterSparseCsrCPU_0.cpp:737 [kernel]\nSparseCsrMeta: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterSparseCsrMeta_0.cpp:702 [kernel]\nBackendSelect: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterBackendSelect.cpp:792 [kernel]\nPython: registered at /workspaces/torch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:194 [backend fallback]\nFuncTorchDynamicLayerBackMode: registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/DynamicLayer.cpp:503 [backend fallback]\nFunctionalize: registered at /workspaces/torch/pytorch/aten/src/ATen/FunctionalizeFallbackKernel.cpp:349 [backend fallback]\nNamed: registered at /workspaces/torch/pytorch/aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]\nConjugate: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/ConjugateFallback.cpp:21 [kernel]\nNegative: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/native/NegateFallback.cpp:22 [kernel]\nZeroTensor: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/ZeroTensorFallback.cpp:90 [kernel]\nADInplaceOrView: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:100 [backend fallback]\nAutogradOther: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradCPU: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradCUDA: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradHIP: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradXLA: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradMPS: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradIPU: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradXPU: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradHPU: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradVE: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradLazy: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradMTIA: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradPrivateUse1: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradPrivateUse2: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradPrivateUse3: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradMeta: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradNestedTensor: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nTracer: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/TraceType_2.cpp:17859 [kernel]\nAutocastCPU: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/autocast_mode.cpp:322 [backend fallback]\nAutocastXPU: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/autocast_mode.cpp:465 [backend fallback]\nAutocastMPS: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/autocast_mode.cpp:209 [backend fallback]\nAutocastCUDA: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/autocast_mode.cpp:165 [backend fallback]\nFuncTorchBatched: registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp:731 [backend fallback]\nBatchedNestedTensor: registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp:758 [backend fallback]\nFuncTorchVmapMode: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/VmapModeRegistrations.cpp:27 [backend fallback]\nBatched: registered at /workspaces/torch/pytorch/aten/src/ATen/LegacyBatchingRegistrations.cpp:1075 [backend fallback]\nVmapMode: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]\nFuncTorchGradWrapper: registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/TensorWrapper.cpp:208 [backend fallback]\nPythonTLSSnapshot: registered at /workspaces/torch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:202 [backend fallback]\nFuncTorchDynamicLayerFrontMode: registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/DynamicLayer.cpp:499 [backend fallback]\nPreDispatch: registered at /workspaces/torch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:206 [backend fallback]\nPythonDispatcher: registered at /workspaces/torch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:198 [backend fallback]\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNotImplementedError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Here's a JAX tensor\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m b \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrandn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mjax\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m      3\u001b[0m b\n",
      "\u001b[0;31mNotImplementedError\u001b[0m: Could not run 'aten::empty.memory_format' with arguments from the 'jax' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'aten::empty.memory_format' is only available for these backends: [CPU, Meta, QuantizedCPU, QuantizedMeta, MkldnnCPU, SparseCPU, SparseMeta, SparseCsrCPU, SparseCsrMeta, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradHIP, AutogradXLA, AutogradMPS, AutogradIPU, AutogradXPU, AutogradHPU, AutogradVE, AutogradLazy, AutogradMTIA, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, AutogradMeta, AutogradNestedTensor, Tracer, AutocastCPU, AutocastXPU, AutocastMPS, AutocastCUDA, FuncTorchBatched, BatchedNestedTensor, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PreDispatch, PythonDispatcher].\n\nCPU: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterCPU_1.cpp:2513 [kernel]\nMeta: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterMeta_0.cpp:5433 [kernel]\nQuantizedCPU: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterQuantizedCPU_0.cpp:311 [kernel]\nQuantizedMeta: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterQuantizedMeta_0.cpp:116 [kernel]\nMkldnnCPU: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterMkldnnCPU_0.cpp:230 [kernel]\nSparseCPU: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterSparseCPU_0.cpp:843 [kernel]\nSparseMeta: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterSparseMeta_0.cpp:187 [kernel]\nSparseCsrCPU: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterSparseCsrCPU_0.cpp:737 [kernel]\nSparseCsrMeta: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterSparseCsrMeta_0.cpp:702 [kernel]\nBackendSelect: registered at /workspaces/torch/pytorch/build/aten/src/ATen/RegisterBackendSelect.cpp:792 [kernel]\nPython: registered at /workspaces/torch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:194 [backend fallback]\nFuncTorchDynamicLayerBackMode: registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/DynamicLayer.cpp:503 [backend fallback]\nFunctionalize: registered at /workspaces/torch/pytorch/aten/src/ATen/FunctionalizeFallbackKernel.cpp:349 [backend fallback]\nNamed: registered at /workspaces/torch/pytorch/aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]\nConjugate: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/ConjugateFallback.cpp:21 [kernel]\nNegative: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/native/NegateFallback.cpp:22 [kernel]\nZeroTensor: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/ZeroTensorFallback.cpp:90 [kernel]\nADInplaceOrView: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:100 [backend fallback]\nAutogradOther: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradCPU: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradCUDA: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradHIP: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradXLA: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradMPS: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradIPU: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradXPU: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradHPU: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradVE: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradLazy: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradMTIA: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradPrivateUse1: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradPrivateUse2: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradPrivateUse3: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradMeta: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nAutogradNestedTensor: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/VariableType_2.cpp:20171 [autograd kernel]\nTracer: registered at /workspaces/torch/pytorch/torch/csrc/autograd/generated/TraceType_2.cpp:17859 [kernel]\nAutocastCPU: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/autocast_mode.cpp:322 [backend fallback]\nAutocastXPU: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/autocast_mode.cpp:465 [backend fallback]\nAutocastMPS: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/autocast_mode.cpp:209 [backend fallback]\nAutocastCUDA: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/autocast_mode.cpp:165 [backend fallback]\nFuncTorchBatched: registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp:731 [backend fallback]\nBatchedNestedTensor: registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp:758 [backend fallback]\nFuncTorchVmapMode: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/VmapModeRegistrations.cpp:27 [backend fallback]\nBatched: registered at /workspaces/torch/pytorch/aten/src/ATen/LegacyBatchingRegistrations.cpp:1075 [backend fallback]\nVmapMode: fallthrough registered at /workspaces/torch/pytorch/aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]\nFuncTorchGradWrapper: registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/TensorWrapper.cpp:208 [backend fallback]\nPythonTLSSnapshot: registered at /workspaces/torch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:202 [backend fallback]\nFuncTorchDynamicLayerFrontMode: registered at /workspaces/torch/pytorch/aten/src/ATen/functorch/DynamicLayer.cpp:499 [backend fallback]\nPreDispatch: registered at /workspaces/torch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:206 [backend fallback]\nPythonDispatcher: registered at /workspaces/torch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:198 [backend fallback]\n"
     ]
    }
   ],
   "source": [
    "# Here's a JAX tensor. But this fails because we did not enable the environment.\n",
    "b = torch.randn(2, 3, device='jax')\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(<class 'jaxlib.xla_extension.ArrayImpl'> [[-0.26367468 -1.1492677   0.04950856]\n",
      " [ 0.22234005 -1.4643176   0.07239702]])\n"
     ]
    }
   ],
   "source": [
    "# We have to run PyTorch code in a torchax environment to use JAX tensors.\n",
    "with torchax.default_env():\n",
    "  b = torch.randn(2, 3, device='jax')\n",
    "  print(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}