python3Packages.mujoco-mjx: init at 3.3.0

Co-authored-by: Gaétan Lepage <33058747+GaetanLepage@users.noreply.github.com>

pkgs/development/python-modules/mujoco-mjx/default.nix

@@ -0,0 +1,50 @@
+{
+  lib,
+  buildPythonPackage,
+
+  # src / metadata
+  mujoco-main,
+
+  # build-system
+  setuptools,
+
+  # dependencies
+  absl-py,
+  etils,
+  importlib-resources,
+  jax,
+  jaxlib,
+  mujoco,
+  scipy,
+  trimesh,
+}:
+
+buildPythonPackage {
+  pname = "mujoco-mjx";
+  inherit (mujoco-main) src version;
+
+  pyproject = true;
+
+  sourceRoot = "${mujoco-main.src.name}/mjx";
+
+  build-system = [ setuptools ];
+
+  dependencies = [
+    absl-py
+    etils
+    importlib-resources
+    jax
+    jaxlib
+    mujoco
+    scipy
+    trimesh
+  ] ++ etils.optional-dependencies.epath;
+
+  pythonImportsCheck = [ "mujoco.mjx" ];
+
+  meta = {
+    description = "MuJoCo XLA (MJX)";
+    inherit (mujoco.meta) homepage changelog license;
+    maintainers = with lib.maintainers; [ nim65s ];
+  };
+}

pkgs/top-level/python-packages.nix

@@ -10123,6 +10123,8 @@ self: super: with self; {
 
   mujoco = callPackage ../development/python-modules/mujoco { inherit (pkgs) mujoco; };
 
+  mujoco-mjx = callPackage ../development/python-modules/mujoco-mjx { mujoco-main = pkgs.mujoco; };
+
   mujson = callPackage ../development/python-modules/mujson { };
 
   mullvad-api = callPackage ../development/python-modules/mullvad-api { };