---
license: apache-2.0
tags:
  - kernel
---



```bash
nix develop --show-trace -i -L .#test --command python -m pytest -s tests
```

expected output:

```
============== test session starts ===============
platform linux -- Python 3.12.10, pytest-8.3.5, pluggy-1.5.0
rootdir: /home/ubuntu/Projects/megablocks-moe
plugins: hypothesis-6.130.12
collecting 43 items                              world_size=1
collected 387 items

tests/layers/moe_test.py ...........................................
tests/ops/binned_gather_test.py .....................
tests/ops/binned_scatter_test.py .....................
tests/ops/cumsum_test.py ................................
tests/ops/histogram_test.py ......................................................
tests/ops/padded_gather_test.py ......................................
tests/ops/padded_scatter_test.py ......................................................
tests/ops/replicate_test.py ..................................................................................
tests/ops/sort_test.py ..................
tests/ops/topology_test.py ....................
tests/test_mb_moe.py megablocks_moe module imported successfully.
Available functions: ['Arguments', 'MLP', 'MoE', 'ParallelDroplessMLP', 'ParallelMLP', 'SparseGLU', 'SparseMLP', '__all__', '__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__path__', '__spec__', '_megablocks_a4f6452_dirty', '_ops', 'argsort', 'backend', 'cumsum', 'dMoE', 'exclusive_cumsum', 'get_load_balancing_loss', 'grouped_gemm_util', 'histogram', 'inclusive_cumsum', 'indices', 'layers', 'ops', 'replicate_backward', 'replicate_forward', 'sort', 'torch']
.cumsum output: tensor([0, 1, 3, 6], device='cuda:0', dtype=torch.int16)
...

================ warnings summary ================
...
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
======= 387 passed, 18 warnings in 54.63s ========
```