Project import generated by Copybara.

GitOrigin-RevId: e2dbd9f87a8c5c077a962d90fa73bd6286e30c89

Author: Copybara

.devcontainer/devcontainer.json

@@ -0,0 +1,52 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the
+// README at: https://github.com/devcontainers/templates/tree/main/src/docker-existing-dockerfile
+{
+	"name": "Jaxpp Latest",
+	"image": "gitlab-master.nvidia.com:5005/cml/jaxpp_dev/jaxpp:main",
+	// "build": {
+	// 	// Sets the run context to one level up instead of the .devcontainer folder.
+	// 	"context": "..",
+	// 	// Update the 'dockerFile' property if you aren't using the standard 'Dockerfile' filename.
+	// 	"dockerfile": "../benchmarks/paxml/Dockerfile",
+	// },
+	"runArgs": [
+		"--gpus=all",
+		"--env=NVIDIA_DISABLE_REQUIRE=1", // NOTE comment out if host has correct cuda driver
+		"--shm-size=10.24gb"
+	],
+	// NOTE(shadow): devcontainer.json doesn't allow shadowing an existing folder
+	// in the container (e.g. /workdir/jaxpp).
+	// While it may seem a good idea to mount it with `-v` in `runArgs` above
+	// that doesn't work with vscode's git integration.
+	"containerUser": "jaxpp-dev",
+	"containerEnv": {"EDITOR": "code -w"},
+	// Features to add to the dev container. More info: https://containers.dev/features.
+	// "features": {},
+	// Use 'forwardPorts' to make a list of ports inside the container available locally.
+	// "forwardPorts": [],
+	// NOTE(shadow): Since we cannot shadow /workdir/jaxpp we update the editable
+	//  version of jaxpp to the workspace one
+	"postCreateCommand": "sudo rm -rf /workdir/jaxpp && pip install -e '/workspaces/${localWorkspaceFolderBasename}[dev]' --no-deps",
+	// Configure tool-specific properties.
+	"customizations": {
+		"vscode": {
+			"extensions": [
+				"ms-python.python",
+				"ms-python.vscode-pylance",
+				"charliermarsh.ruff",
+				"tamasfe.even-better-toml"
+			],
+			"settings": {
+				"terminal.integrated.defaultProfile.linux": "bash",
+				"terminal.integrated.profiles.linux": {
+					"zsh": {
+						"path": "/bin/bash"
+					}
+				},
+				"python.defaultInterpreterPath": "/home/jaxpp-dev/dev-env/bin/python3"
+			}
+		}
+	}
+	// Uncomment to connect as an existing user other than the container default. More info: https://aka.ms/dev-containers-non-root.
+	// "remoteUser": "devcontainer"
+}

.github/ISSUE_TEMPLATE.md

@@ -0,0 +1,19 @@
+* JaxPP version:
+* Jax version:
+* cuDNN version:
+* cuBLAS version:
+* NCCL version:
+* Python version:
+* Operating System:
+
+### Description
+
+Describe what you were trying to get done.
+Tell us what happened, what went wrong, and what you expected to happen.
+
+### What I Did
+
+```
+Paste the command(s) you ran and the output.
+If there was a crash, please include the traceback here.
+```

.vscode/launch.json

@@ -0,0 +1,22 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "basic.py",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/benchmarks/basic.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "args": [
+                "--dtype=float16"
+            ],
+            "env": {
+                "RAY_ADDRESS": "local"
+            }
+        }
+    ]
+}

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.analysis.typeCheckingMode": "basic"
+}