meta.yaml

{% set version = "0.3.14" %}
{% set name = "jaxlib" %}

{% set number = 3 %}

{% if cuda_compiler_version != "None" %}
{% set number = number + 200 %}
{% endif %}


package:
  name: jaxlib
  version: {{ version }}

source:
  # only pull sources after upstream PyPI release...
  url: https://github.com/google/jax/archive/{{ name }}-v{{ version }}.tar.gz
  sha256: 6d6b660cb915b3fbf151dc67a471623e5f83b6a66493b9e438273d290a43b40f
  patches:
    - 0001-Don-t-add-dashes-to-version.patch
    - 0001-Patch-pocketfft-for-C-17-osx-64.patch  # [osx and x86_64]
    - 0002-bazelrc-112.diff  # [cuda_compiler_version == "11.2"]

build:
  number: {{ number }}
  skip: true  # [win or py<37]
  skip: true  # [cuda_compiler_version == "10.2"]
  skip: true  # [cuda_compiler_version == "11.0"]
  skip: true  # [cuda_compiler_version == "11.1"]
  skip: true  # [abseil_cpp != '20211102.0']
  string: cuda{{ cuda_compiler_version | replace('.', '') }}py{{ CONDA_PY }}h{{ PKG_HASH }}_{{ PKG_BUILDNUM }}  # [cuda_compiler_version != "None"]
  string: cpu_py{{ CONDA_PY }}h{{ PKG_HASH }}_{{ PKG_BUILDNUM }}  # [cuda_compiler_version == "None"]

requirements:
  build:
    - {{ compiler('c') }}
    - {{ compiler('cxx') }}
    - {{ compiler('cuda') }}                 # [cuda_compiler_version != "None"]
    - python                                 # [build_platform != target_platform]
    - cross-python_{{ target_platform }}     # [build_platform != target_platform]
    - numpy                                  # [build_platform != target_platform]
    - scipy                                  # [build_platform != target_platform]
    - unzip
    # Keep bazel listed twice here to help the migrators track dependencies
    - bazel
    - bazel  >=5.1.1
    - bazel-toolchain
    # need protoc
    - libprotobuf
    # needs protoc-gen-grpc
    - grpc-cpp
    # needs flatc
    - flatbuffers <2.0.6
  host:
    - cudnn      # [cuda_compiler_version != "None"]
    - nccl       # [cuda_compiler_version != "None"]
    - python
    - pip
    - numpy
    - scipy
    - wheel
    # Reactivate once we have switched the global abseil pin to a new version
    # - abseil-cpp
    - flatbuffers <2.0.6
    - grpc-cpp
    - openssl
    - zlib
  run:
    - python
    - absl-py
    - {{ pin_compatible('numpy') }}
    - scipy >=1.5
    - python-flatbuffers >=1.12,<3.0
    - __cuda  # [cuda_compiler_version != "None"]
  run_constrained:
    - jax >={{ version }}
    - __osx >={{ MACOSX_DEPLOYMENT_TARGET|default("10.9") }}  # [osx]

test:
  files:
    - test_jaxlib.py
  requires:
    - pip
  commands:
    - pip check
    - python test_jaxlib.py
  imports:
    - jaxlib
    # Only this import actually triggers the load of XLA
    - jaxlib.xla_client

about:
  home: http://github.com/google/jax
  license: Apache-2.0
  license_file: LICENSE
  summary: 'Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more'
  dev_url: http://github.com/google/jax

extra:
  recipe-maintainers:
    - ericmjl
    - xhochy
    - ngam