# Member of Technical Staff, Agentic Environments

**Company:** [Cohere](https://hotfix.jobs/companies/cohere)
**Location:** Remote
**Role:** ML Engineering
**Skills:** Python, JAX, PyTorch, xla, mlir, Kubernetes, Ray, Distributed Training, memory profiling, speed profiling, LLMs, Machine Learning, test-driven development
**Posted:** 2026-08-13

> Build scalable software and tools for frontier model training, research experimentation, and production machine learning systems. The role requires strong Python and distributed-training expertise, experience with ML frameworks and infrastructure, and the ability to optimize and debug large language model systems.

## Job Description

## Responsibilities
- Design and write high-performing, scalable software for training models.
- Develop tools to support and accelerate research and large language model training.
- Coordinate with Infrastructure, Efficiency, Serving, Agent, Multimodal, and Multilingual teams.
- Research, implement, and experiment with ideas on cluster and data infrastructure.
- Collaborate with scientists, engineers, and cross-functional teams.

## Requirements
- Extremely strong software engineering skills.
- Commitment to test-driven development, clean code, and reducing technical debt.
- Proficiency in Python and machine learning frameworks such as JAX, PyTorch, and/or XLA/MLIR.
- Experience using and debugging large-scale distributed training strategies, including memory and speed profiling.
- Experience with distributed training infrastructure such as Kubernetes and associated frameworks such as Ray.
- Experience in machine learning and large language model academic research.
- Ability to tune and optimize large language models.
- Ability to navigate complex machine learning codebases and resolve issues.
- Ability to work effectively with colleagues across experience levels in fast-paced, technically challenging environments.

## Compensation and Benefits
- Weekly lunch stipend of $75/£75 or equivalent in local currency.
- Full health and dental benefits, including a separate mental health budget.
- RRSP matching, 401(k), or pension scheme.
- 100% parental leave top-up for up to six months for either parent.
- Annual enrichment benefits for arts and culture, fitness and wellness, quality time, and workspace improvements.
- Education and learning stipend for conferences, courses, and coaching.
- Six weeks of paid vacation (30 working days).
- Travel budget for remote employees visiting other offices and an annual company offsite.
- Coworking benefit for employees not near an office.
- $500 home office stipend.

## Similar roles

- [Staff Machine Learning Engineer, Traffic Intelligence](https://hotfix.jobs/jobs/cee28c3a-eb64-42f7-a584-8d2b4a0980b8) - Airbnb - Remote - $212k – $265k/yr
- [Staff Machine Learning Infrastructure Engineer, Embedding Platform](https://hotfix.jobs/jobs/9b4387bf-ac16-48ba-810b-2e6a612e3617) - Reddit - Remote - $253k – $355k/yr
- [Senior Staff Machine Learning Engineer, Feed Relevance](https://hotfix.jobs/jobs/a2f9b178-c6cf-49bf-ba93-00688d1027fb) - Reddit - Remote - $266k – $372k/yr
- [Staff Applied Scientist](https://hotfix.jobs/jobs/d518c7a7-07d4-4f06-af79-9c478df769ca) - Garner Health - New York, NY - $300k – $390k/yr
- [Senior/Staff Software Engineer, ML Inference Platform](https://hotfix.jobs/jobs/1b014e68-d8bc-4d84-b03a-d94b7a57c90f) - Nuro - Mountain View, CA - $194k – $352k/yr

**Apply:** https://hotfix.jobs/jobs/f3ceb248-2d3e-4b28-a5d1-10891969916e
**Canonical:** https://hotfix.jobs/jobs/f3ceb248-2d3e-4b28-a5d1-10891969916e