Brutax

When all else fails, why not a brute-force search! Brutax is a JAX library for function optimization by brute-force grid search. Features include:

Highly-parallel function evaluations
PyTree-valued grids
Customizable search behavior downstream
Smooth integration with JAX function transformations: JIT, autodiff, vectorization, and scaling across GPU/TPUs

Installation

pip install brutax

Quick example

import brutax
import jax.numpy as jnp

def fn(tree_grid_point, _):
    x, y = tree_grid_point
    return (x - 1.)**2 + (y - 2.)**2

# The `tree_grid` is the cartesian product of x and y grids
x_grid, y_grid = (jnp.arange(-5., 5., 1), jnp.arange(-5., 5., 0.1))
tree_grid = (x_grid, y_grid)
# Run grid search over (10 x 100) grid
sol = brutax.run_grid_search(
    fn, method=brutax.MinimumSearchMethod(), tree_grid=tree_grid, args=None
)
x_min, y_min = sol.value
assert jnp.isclose(x_min, 1.) and jnp.isclose(y_min, 2.)

Acknowledgements

The design of brutax is heavily inspired from the JAX non-linear optimization library optimistix.

Development Status
- 4 - Beta
Intended Audience
- Developers
- Science/Research
License
- OSI Approved :: Apache Software License
Operating System
- OS Independent
Programming Language
- Python
- Python :: 3

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

brutax-0.1.0.tar.gz (18.6 kB view details)

Uploaded Aug 26, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

brutax-0.1.0-py3-none-any.whl (19.8 kB view details)

Uploaded Aug 26, 2026 Python 3

File details

Details for the file brutax-0.1.0.tar.gz.

File metadata

Download URL: brutax-0.1.0.tar.gz
Upload date: Aug 26, 2026
Size: 18.6 kB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for brutax-0.1.0.tar.gz
Algorithm	Hash digest
SHA256	`9f4889c0543dd82123a3345995891d5fdf8fe34e06efa03a5b1b9799066282f4`
MD5	`4f05245f15cda6b9665d3e60a9805ffb`
BLAKE2b-256	`aaea1ab4fd741ededd300db9f038d33b1bf3063da08afd774831425c32f2cb46`

See more details on using hashes here.

Provenance

The following attestation bundles were made for brutax-0.1.0.tar.gz:

Publisher: publish.yml on michael-0brien/brutax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Statement:
- Statement type: https://in-toto.io/Statement/v1
- Predicate type: https://docs.pypi.org/attestations/publish/v1
- Subject name: brutax-0.1.0.tar.gz
- Subject digest: 9f4889c0543dd82123a3345995891d5fdf8fe34e06efa03a5b1b9799066282f4
- Sigstore transparency entry: 2607163790
- Sigstore integration time: Aug 26, 2026
Source repository:
- Permalink: michael-0brien/brutax@da628aa41c661d1c8bbc782ff5af9492b2421fdf
- Branch / Tag: refs/tags/v0.1.0
- Owner: https://github.com/michael-0brien
- Access: public
Publication detail:
- Token Issuer: https://token.actions.githubusercontent.com
- Runner Environment: github-hosted
- Publication workflow: publish.yml@da628aa41c661d1c8bbc782ff5af9492b2421fdf
- Trigger Event: release