refactor: clean up lax.broadcast_in_dim_p

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 96.59%. Comparing base (0022969) to head (4b10079). Report is 1 commits behind head on main.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #204 +/- ## ======================================= Coverage 96.59% 96.59% ======================================= Files 37 37 Lines 1498 1498 ======================================= Hits 1447 1447 Misses 51 51 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.