From 9a6297dc5e8197d53fe7aa02ee847babdbfbb628 Mon Sep 17 00:00:00 2001
From: Guilhem Saurel <guilhem.saurel@laas.fr>
Date: Wed, 26 Mar 2025 13:53:16 +0100
Subject: [PATCH] python3Packages.mujoco-mjx: init at 3.3.0
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Gaétan Lepage <33058747+GaetanLepage@users.noreply.github.com>
---
 .../python-modules/mujoco-mjx/default.nix     | 50 +++++++++++++++++++
 pkgs/top-level/python-packages.nix            |  2 +
 2 files changed, 52 insertions(+)
 create mode 100644 pkgs/development/python-modules/mujoco-mjx/default.nix

diff --git a/pkgs/development/python-modules/mujoco-mjx/default.nix b/pkgs/development/python-modules/mujoco-mjx/default.nix
new file mode 100644
index 000000000000..7073dc4484d1
--- /dev/null
+++ b/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 ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 01db0f0e07f1..704825453925 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/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 { };