{"id":13478372,"url":"https://github.com/lucidrains/jax2torch","last_synced_at":"2025-04-06T02:09:46.721Z","repository":{"id":62572193,"uuid":"421224447","full_name":"lucidrains/jax2torch","owner":"lucidrains","description":"Use Jax functions in Pytorch","archived":false,"fork":false,"pushed_at":"2023-07-01T17:24:32.000Z","size":18,"stargazers_count":239,"open_issues_count":1,"forks_count":9,"subscribers_count":5,"default_branch":"main","last_synced_at":"2025-03-30T01:11:17.945Z","etag":null,"topics":["deep-learning-framework","jax","torch"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2021-10-26T00:16:32.000Z","updated_at":"2025-03-24T02:16:41.000Z","dependencies_parsed_at":"2024-10-23T01:32:25.728Z","dependency_job_id":null,"html_url":"https://github.com/lucidrains/jax2torch","commit_stats":{"total_commits":17,"total_committers":2,"mean_commits":8.5,"dds":"0.17647058823529416","last_synced_commit":"cd6c38a47827c21b6568ddd4f345c4ed95797350"},"previous_names":[],"tags_count":8,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fjax2torch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fjax2torch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fjax2torch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fjax2torch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/jax2torch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247423515,"owners_count":20936626,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning-framework","jax","torch"],"created_at":"2024-07-31T16:01:56.139Z","updated_at":"2025-04-06T02:09:46.696Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"## jax2torch\n\nUse Jax functions in Pytorch with DLPack, as outlined \u003ca href=\"https://gist.github.com/mattjj/e8b51074fed081d765d2f3ff90edf0e9\"\u003ein a gist\u003c/a\u003e by \u003ca href=\"https://github.com/mattjj\"\u003e@mattjj\u003c/a\u003e. The repository was made for the purposes of making this \u003ca href=\"https://github.com/spetti/SMURF\"\u003edifferentiable alignment work\u003c/a\u003e interoperable with Pytorch projects.\n\n## Install\n\n```bash\n$ pip install jax2torch\n```\n\n## Memory management\n\nBy default, Jax pre-allocates 90% of VRAM, which leaves Pytorch with very little left over.  To prevent this behavior, set the `XLA_PYTHON_CLIENT_PREALLOCATE` environmental variable to false before running any Jax code:\n\n```python\nimport os\nos.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"false\"\n```\n\n## Usage\n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1GBEEnpuCvLS1bhb_xGCO5Y40rFiQrh6G?usp=sharing) Quick test\n\n```python\nimport jax\nimport torch\nfrom jax2torch import jax2torch\nimport os\n\nos.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"false\"\n\n# Jax function\n\n@jax.jit\ndef jax_pow(x, y = 2):\n  return x ** y\n\n# convert to Torch function\n\ntorch_pow = jax2torch(jax_pow)\n\n# run it on Torch data!\n\nx = torch.tensor([1., 2., 3.])\ny = torch_pow(x, y = 3)\nprint(y)  # tensor([1., 8., 27.])\n\n# And differentiate!\n\nx = torch.tensor([2., 3.], requires_grad = True)\ny = torch.sum(torch_pow(x, y = 3))\ny.backward()\nprint(x.grad) # tensor([12., 27.])\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fjax2torch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fjax2torch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fjax2torch/lists"}