Importing functions from max.nn

[avb03]

I am referencing some of the library functions under the max/nn directory on the modular github.

For example I am trying to import:

from nn.mha import _naive_attention_with_transpose, mha_gpu_naive
from nn.mha_mask import CausalMask, MaterializedMask

as one of the benchmark scripts does. I was able to get this script running with this import a couple weeks ago, but I had to reinstall the pixi shell with modular today and am getting these errors:

mojo: error: failed to parse the provided Mojo source module

wherever I try to import from nn. How can I import these functions?

[Ehsan (Ehsan M. Kermani)]

For those Mojo modules, are you following this guide to build and import? I don’t think pixi can work with that as is.

[avb03]

Thanks for the prompt reply. I get this error when trying to follow the steps outlined in that post. When running ./bazelw build //max/kernels/src/nn, I get this error:

//mojo/stdlib/stdlib:stdlib building mojo package failed: (Exit 1): mojo failed: error executing MojoPackage command (from target //mojo/stdlib/stdlib:stdlib) external/rules_mojo++mojo+mojo_toolchain_linux_x86_64/bin/mojo package '-strip-file-prefix=.' -o bazel-out/k8-fastbuild/bin/mojo/stdlib/stdlib/stdlib.mojopkg mojo/stdlib/stdlib