python312Packages.tensorflow: cleanup dependencies (#391992)

{

  stdenv,

  lib,

  fetchurl,

  stdenv,

  buildPythonPackage,

  isPy3k,

  astor,

  fetchurl,

  # buildInputs

  llvmPackages,

  # build-system

  distutils,

  # dependencies

  ml-dtypes,

  absl-py,

  astunparse,

  flatbuffers,

  gast,

  google-pasta,

  wrapt,

  grpcio,

  h5py,

  libclang,

  numpy,

  six,

  termcolor,

  opt-einsum,

  packaging,

  protobuf,

  absl-py,

  grpcio,

  mock,

  scipy,

  distutils,

  wheel,

  jax,

  ml-dtypes,

  opt-einsum,

  tensorflow-estimator-bin,

  requests,

  six,

  tensorboard,

  termcolor,

  typing-extensions,

  wrapt,

  isPy3k,

  mock,

  config,

  cudaSupport ? config.cudaSupport,

  cudaPackages,

  zlib,

  python,

  addDriverRunpath,

  astunparse,

  flatbuffers,

  h5py,

  llvmPackages,

  typing-extensions,

}:

# We keep this binary build for three reasons:

  inherit (cudaPackages) cudatoolkit cudnn;

  isCudaJetson = cudaSupport && cudaPackages.cudaFlags.isJetsonBuild;

  isCudaX64 = cudaSupport && stdenv.hostPlatform.isx86_64;

in

buildPythonPackage rec {

  pname = "tensorflow" + lib.optionalString cudaSupport "-gpu";

  buildInputs = [ llvmPackages.openmp ];

  build-system = [

    distutils

  ];

  nativeBuildInputs =

    lib.optionals cudaSupport [ addDriverRunpath ]

    ++ lib.optionals isCudaJetson [ cudaPackages.autoAddCudaCompatRunpath ];

  dependencies = [

    absl-py

    astunparse

    flatbuffers

    typing-extensions

    distutils

    gast

    google-pasta

    grpcio

    h5py

    libclang

    ml-dtypes

    numpy

    opt-einsum

    packaging

    protobuf

    numpy

    scipy

    (if isCudaX64 then jax else ml-dtypes)

    termcolor

    grpcio

    requests

    six

    astor

    absl-py

    gast

    opt-einsum

    google-pasta

    wrapt

    tensorflow-estimator-bin

    tensorboard

    h5py

    termcolor

    typing-extensions

    wrapt

  ] ++ lib.optional (!isPy3k) mock;

  build-system =

    [

      distutils

      wheel

    ]

    ++ lib.optionals cudaSupport [ addDriverRunpath ]

    ++ lib.optionals isCudaJetson [ cudaPackages.autoAddCudaCompatRunpath ];

  preConfigure = ''

    unset SOURCE_DATE_EPOCH

      mv "$f" "$(sed -E 's/(nv[0-9]+)\.0*([0-9]+)/\1.\2/' <<< "$f")"

    done

    wheel unpack --dest unpacked ./*.whl

    rm ./*.whl

    (

      cd unpacked/tensorflow*

      # Adjust dependency requirements:

      # - Relax flatbuffers, gast, protobuf, tensorboard, and tensorflow-estimator version requirements that don't match what we have packaged

      # - The purpose of python3Packages.libclang is not clear at the moment and we don't have it packaged yet

      # - keras and tensorlow-io-gcs-filesystem will be considered as optional for now.

      # - numpy was pinned to fix some internal tests: https://github.com/tensorflow/tensorflow/issues/60216

      sed -i *.dist-info/METADATA \

        -e "/Requires-Dist: flatbuffers/d" \

        -e "/Requires-Dist: gast/d" \

        -e "/Requires-Dist: keras/d" \

        -e "/Requires-Dist: libclang/d" \

        -e "/Requires-Dist: protobuf/d" \

        -e "/Requires-Dist: tensorboard/d" \

        -e "/Requires-Dist: tensorflow-estimator/d" \

        -e "/Requires-Dist: tensorflow-io-gcs-filesystem/d" \

        -e "s/Requires-Dist: numpy (.*)/Requires-Dist: numpy/"

    )

    wheel pack ./unpacked/tensorflow*

    popd

  '';