{"id":13564313,"url":"https://github.com/srush/parallax","last_synced_at":"2025-09-13T16:26:27.128Z","repository":{"id":66268224,"uuid":"265112884","full_name":"srush/parallax","owner":"srush","description":null,"archived":false,"fork":false,"pushed_at":"2020-05-25T18:24:30.000Z","size":66,"stargazers_count":153,"open_issues_count":1,"forks_count":4,"subscribers_count":5,"default_branch":"master","last_synced_at":"2025-04-12T15:12:33.281Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/srush.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2020-05-19T01:49:51.000Z","updated_at":"2025-01-09T23:57:30.000Z","dependencies_parsed_at":"2023-06-13T04:15:18.410Z","dependency_job_id":null,"html_url":"https://github.com/srush/parallax","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/srush%2Fparallax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/srush%2Fparallax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/srush%2Fparallax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/srush%2Fparallax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/srush","download_url":"https://codeload.github.com/srush/parallax/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248586245,"owners_count":21128998,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-08-01T13:01:29.567Z","updated_at":"2025-04-12T15:12:38.773Z","avatar_url":"https://github.com/srush.png","language":"Python","funding_links":[],"categories":["Python","Libraries"],"sub_categories":["Inactive Libraries","New Libraries"],"readme":"# Parallax - Immutable Torch Modules for JAX\n\n\n\u003cimg width=450px src=\"https://developers.google.com/web/updates/images/2016/12/performant-parallaxing/parallax.jpg\"\u003e\n\n\nParallax is a prototype for a pure module system for JAX implemented by Sabrina Mielke (@sjmielke) and Sasha Rush (@srush).\n\nMain ideas:\n\n* Make param modules immutable trees.\n* Replace all imperative style coding and init.\n* Avoid tracking state for most applications by first distributing seeds / globals through tree.\n\n```python\n\nfrom parallax import Module, Parameter, ParamInit\n\nclass Dense(Module):\n    # All parameter-holders are explicitly declared.\n    weight : Parameter\n    bias : Parameter\n\n    # Setup replace __init__ and creates shapes and binds lazy initializers.\n    def __init__(self, in_size, out_size):\n        super().__init__()\n        self.weight = ParamInit((out_size, in_size), init.xavier_normal())\n        self.bias = ParamInit((out_size,), init.normal())\n\n\n    # Forward is just like standard pytorch.\n    def forward(self, input):\n        return self.weight @ input + self.bias\n\n    # Hook for pretty printing\n    def extra_repr(self):\n        return \"%d, %d\"%(self.weight.shape[1], self.weight.shape[0])\n\nclass Dropout(Module):\n    # Arbitrary constants allowed.\n    rate : float\n    def __init__(self, rate):\n        super().__init__()\n        self.rate = rate\n\n    def forward(self, input):\n        # RNG state is use-once or split. Attached to tree.\n        state = self.rng\n\n        if self.mode == \"train\":\n            keep = jax.random.bernoulli(state, self.rate, input.shape)\n            return jax.numpy.where(keep, input / self.rate, 0)\n        else:\n            return input\n\nclass BinaryNetwork(Module):\n    # No difference between modules and parameters\n    dense1 : Dense\n    dense2 : Dense\n    dense3 : Dense\n    dropout : Dropout\n\n    def __init__(self, input_size, hidden_size):\n        super().__init__()\n        self.dense1 = Dense(input_size, hidden_size)\n        self.dense2 = Dense(hidden_size, hidden_size)\n        self.dense3 = Dense(hidden_size, 1)\n        self.dropout = Dropout(0.2)\n\n    def forward(self, input):\n\n        # Standard usage works out of the box.\n        x = jax.numpy.tanh(self.dense1(input))\n\n        # Stochastic modules (have random seed already)\n        x = self.dropout(x)\n\n        # Shared params / recurrence only requires split to change RNG\n        x = jax.numpy.tanh(self.dense2(x))\n        x = jax.numpy.tanh(self.dense2(x))\n\n        return jax.nn.sigmoid(self.dense3(jax.numpy.tanh(x)))[0]\n\n# Setup param tree -\u003e declarative, immutable\nlayer = BinaryNetwork(5, 10)\nprint(layer)\nprint(layer.dense1)\n\n# Initialize parameters -\u003e stateful, hidden\nrng = jax.random.PRNGKey(0)\nlayer = layer.initialized(rng)\nprint(layer)\nprint(layer.dense1)\n\ninitial_loss = None\nfor i in range(10):\n    # Thread state through parameters -\u003e functor, hidden\n    rng, iter_rng = jax.random.split(rng)\n    layer = layer.new_state(iter_rng, mode=\"train\")\n    \n    # Jax style grad compute -\u003e tree-shaped immutable\n    x = jax.numpy.zeros(5)\n    loss = layer(x)\n    if initial_loss is None:\n        initial_loss = loss\n    print(loss)\n    grad = layer.grad(x)\n    \n    # Grad Update -\u003e tree-shaped\n    layer = jax.tree_util.tree_multimap(lambda p, g: p - 0.3 * g, layer, grad)\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsrush%2Fparallax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsrush%2Fparallax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsrush%2Fparallax/lists"}