pyproject.toml

[tool.poetry]
name = "torch-jax-interop"
version = "0.0.4.post7.dev0+c6122c5"
description = "Utility to convert Tensors from Jax to Torch and vice-versa"
authors = ["Fabrice Normandin <normandf@mila.quebec>"]
license = "MIT"
readme = "README.md"

[tool.poetry.dependencies]
python = "^3.10"
jax = "^0.4.28"
torch = "^2.0.0"
pytorch2jax = "^0.1.0"
flax = "^0.8.4"

[tool.poetry.group.dev.dependencies]
pytest = "^8.2.1"
pytest-regressions = "^2.5.0"
pytest-benchmark = "^4.0.0"
pytest-skip-slow = "^0.0.5"
pre-commit = "^3.7.1"
pytest-testmon = "^2.1.1"
pytest-env = "^1.1.3"
tensor-regression = "^0.0.4"
jax = {extras = ["cuda12"], version = "^0.4.28"}


[tool.poetry-dynamic-versioning]
enable = true
[build-system]
requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"]
build-backend = "poetry_dynamic_versioning.backend"


[tool.pytest_env]
CUBLAS_WORKSPACE_CONFIG=":4096:8"



# pyproject.toml
[tool.pytest.ini_options]
addopts = "--doctest-modules"