Google JAX Jobs in the UK

1 to 25 of 49 Google JAX Jobs in the UK

Machine Learning Engineer

London, United Kingdom
Hybrid / WFH Options
Locus Robotics
MAS). Advanced Python programming skills, with a strong emphasis on writing efficient, scalable, and maintainable code. Proven experience with TensorFlow/PyTorch/Jax, Scikit-learn, and MLOps workflows for training, deployment, and monitoring of ML models. Experience working with Polars and/or Pandas for high-performance More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

Machine Learning Software Engineer, Research

London, UK
Hybrid / WFH Options
PhysicsX Ltd
scaler platforms, e.g., AWS, Azure, GCP); building machine learning models and pipelines in Python, using common libraries and frameworks (e.g., NumPy, SciPy, Pandas, PyTorch, JAX), especially including deep learning applications; C/C++ for computer vision, geometry processing, or scientific computing; software engineering concepts and best practices (e.g., versioning More ❯
Posted:

Software Engineer (GenAI)

London, United Kingdom
Barlowe LLP
or a related field, or equivalent professional experience Strong programming skills in Python, with experience in machine learning frameworks such as PyTorch, TensorFlow or JAX Hands-on experience fine-tuning LLMs and optimising hyperparameters for improved performance Skilled in building LLM-powered applications for real-time decision-making and More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

AI Research Scientist (AEC)

London, United Kingdom
Merantix
Familiarity in statistical methods for Machine Learning (e.g. Bayesian methods, HMMs, graphical models, dimension reduction, clustering, classification, regression techniques, etc.) Familiarity with PyTorch, TensorFlow, JAX or similar frameworks Strong coding abilities in Python Preferred Qualifications Experience in the Architecture, Engineering, and/or Construction domains, including expertise with industry More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

AI Researcher Scientist

London, United Kingdom
Autodesk, Inc
regression techniques, etc.) Familiarity with Deep Learning techniques (e.g. Network architectures, regularization techniques, learning techniques, loss-functions, optimization strategies, etc.) Familiarity with PyTorch , TensorFlow , JAX or similar frameworks Strong coding abilities in Python and/or C++ Preferred Qualifications: 2D & 3D Generative AI Reinforcement Learning LLMs and Natural Language More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

Research Data Scientist

London, United Kingdom
Hybrid / WFH Options
PhysicsX Ltd
data-driven role, with exposure to: building machine learning models and pipelines in Python, using common libraries and frameworks (e.g., NumPy, SciPy, Pandas, PyTorch, JAX), especially including deep learning applications; developing models for bespoke problem settings that involve high-dimensional data (spatiotemporal, geometric, physical); iterating on network architectures and More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

AI Research Engineer

London, UK
Hybrid / WFH Options
Arondite
language understanding The ability to write clean and maintainable code in Python, Rust, C++ or Java Experience with frameworks such as PyTorch, Tensorflow or JAX Strong written and verbal communication skills to document your work effectively and present it to relevant stake-holders and users A passion for identifying More ❯
Posted:

Machine Learning Research Engineer / Scientist

England, UK
Silurian
expect you to have Proven track record in developing and deploying deep learning models. Advanced proficiency in Python and modern ML frameworks (PyTorch, TensorFlow, Jax). Demonstrated experience with distributed training systems and large-scale data pipelines. Strong software engineering practices and system design principles. Excellent problem-solving and More ❯
Posted:

Research Scientist, FAIR (PhD)

London, UK
Meta
and algorithm scaling, or neural network architectures. Experience communicating complex research for public audiences of peers. Experience with deep learning frameworks such as Pytorch, Jax, or Tensorflow. Must obtain work authorization in country of employment at the time of hire, and maintain ongoing work authorization during employment. Preferred Qualifications More ❯
Posted:

AI Research Scientist

Cambridge, Cambridgeshire, United Kingdom
iO Associates
and robotics Key Requirements: PhD in Computer Science, Electrical Engineering, or related field (or equivalent research experience) Strong expertise in ML frameworks (PyTorch, TensorFlow, JAX) and efficient ML techniques (quantization, pruning, distillation, etc.) Experience in on-device ML deployment (mobile CPU, GPU, NPU) and distributed multi-GPU training Fluency More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

Software Engineer- AI/ML, AWS Neuron Distributed Training

London, United Kingdom
Amazon
control management, build processes, testing, and operations experience - Bachelor's degree in computer science or equivalent - Preferred previous software engineer expertise with Pytorch/Jax/Tensorflow, Distributed libraries and Frameworks, End-to-end Model Training. The group presents a lot of opportunity for optimization and scaling large deep More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

Presales System Engineer - AI/HPC - (UK Based EMEA Role)

London, UK
Pure Storage
industry, including knowledge of the landscape and Pure’s solutions with a focus on the following: Deep understanding of industry AI frameworks (pytorch, tensorflow, jax etc) Deep understanding on high performance, low latency, highly parallelized compute solutions at scale (compute, network, storage) Deep understanding of accelerated compute environments (GPU More ❯
Posted:

Machine Learning Scientist

Peterborough, Cambridgeshire, United Kingdom
Hyper Recruitment Solutions Ltd
machine learning, including deep learning and/or generative models. A working knowledge and practical experience with modern ML frameworks (e.g., PyTorch, TensorFlow, or JAX). Key Words: Machine Learning/Drug Discovery/Computational Chemistry/Structural Biology/AI Techniques/Protein-Ligand Interaction/Molecular Modelling More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

Research Scientist

Peterborough, Cambridgeshire, United Kingdom
microTECH Global Limited
experience in academia or industry • We will consider various levels of experience in relevant research areas Key Skills: • Experience with ML frameworks (PyTorch, TensorFlow, JAX) and efficient ML (incl. quantization, pruning, sparsification, distillation, etc.) • Experience with deployment on embedded/mobile devices (such as smartphones, with mobile CPU, GPU More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

Member of Technical Staff, ML Performance

London, UK
Odyssey
of owning projects end to end Problem-solving mindset with the ability to acquire new skills as needed Proficiency with PyTorch (or TF/JAX), as well as NVIDIA GPU ecosystems and optimization stacks Highly metric-based Strong Python and C++ skills Bonus Qualifications Experience optimizing kernels with Triton More ❯
Posted:

Principal Solutions Architect, Generative AI, Startups, AWS

London, United Kingdom
ENGINEERINGUK
and deploying GenAI models and applications on cloud infrastructure. Sound knowledge of machine learning fundamentals, with working knowledge of frameworks such as Pytorch, TensorFlow, JAX or MXNet. Knowledge of MLOps tools and workflows for model development, validation, and deployment. High level of communicating effectively across internal and external organizations More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

Senior Solutions Architect, Generative AI, Startups, AWS

London, UK
Amazon
and deploying GenAI models and applications on cloud infrastructure. - Sound knowledge of machine learning fundamentals, with working knowledge of frameworks such as Pytorch, TensorFlow, JAX or MXNet. - Knowledge of MLOps tools and workflows for model development, validation, and deployment. - High level of communicating effectively across internal and external organizations More ❯
Posted:

Principal Solutions Architect, Generative AI, Startups, AWS

London, United Kingdom
Amazon
and deploying GenAI models and applications on cloud infrastructure. - Sound knowledge of machine learning fundamentals, with working knowledge of frameworks such as Pytorch, TensorFlow, JAX or MXNet. - Knowledge of MLOps tools and workflows for model development, validation, and deployment. - High level of communicating effectively across internal and external organizations More ❯
Employment Type: Permanent
Salary: GBP Annual
Posted:

Machine Learning Scientist

London, UK
Scindo
desire to work on the next generation of bioinformatics and biocatalysis predictive tools. Background in Python and experience with PyTorch; TensorFlow and/or JAX is a bonus, as well as knowledge of other programming languages Data preparation: feature selection, engineering, and data cleaning ML algorithm development and knowledge More ❯
Posted:

Machine Learning Engineer- World-Leading Prop Trading Fund - Oxford Knight

London, UK
Jobs via eFinancialCareers
techniques, model architecture, optimizer, learning rate schedules, etc. Thorough understanding of Python tools and libraries, keen to offer advice on best practices Experience using Jax, Tensorflow or similar ML frameworks a huge plus Benefits: Market-leading salaries Generous benefits package, including physical & mental health benefits, excellent holiday entitlement, significant More ❯
Posted:

Machine Learning Engineer- World-Leading Prop Trading Fund - Oxford Knight

London, UK
Jobs via eFinancialCareers
techniques, model architecture, optimizer, learning rate schedules, etc. Thorough understanding of Python tools and libraries, keen to offer advice on best practices Experience using Jax, Tensorflow or similar ML frameworks a huge plus Benefits: Market-leading salaries Generous benefits package, including physical & mental health benefits, excellent holiday entitlement, significant More ❯
Posted:

Edge AI Experience Kit Lead

Cambridge, Cambridgeshire, UK
Arm Limited
of the latest trends in AI/ML research and development Good familiarity with one of the popular ML frameworks such as PyTorch, TensorFlow, JAX, ONNX, or PaddlePaddle A history of developer engagement or participation in an active open-source development community Experience in Product Owner roles, or a More ❯
Posted:

Machine Learning Engineer- World-Leading Prop Trading Fund | London, UK

London, UK
Oxford Knight
techniques, model architecture, optimizer, learning rate schedules, etc. Thorough understanding of Python tools and libraries, keen to offer advice on best practices Experience using Jax, Tensorflow or similar ML frameworks a huge plus Benefits: Market-leading salaries Generous benefits package, including physical & mental health benefits, excellent holiday entitlement, significant More ❯
Posted:

Ai Software Engineer

United Kingdom
Hybrid / WFH Options
IC Resources
position? Strong C++ development skills Experience with parallel programming - CUDA, OpenCL or SYCL Strong understanding of computer architecture. GPU/CPU Experience with TensorFlow, JAX, NumPy or PyTorch If you are an Ai Software Engineer looking for a new opportunity within a rapidly growing, disruptive start up, please apply More ❯
Posted:

Ai Software Engineer

United Kingdom, UK
Hybrid / WFH Options
JR United Kingdom
this Ai Software Engineer position? Experience with parallel programming - CUDA, OpenCL or SYCL Strong understanding of computer architecture. GPU/CPU Experience with TensorFlow, JAX, NumPy or PyTorch If you are an Ai Software Engineer looking for a new opportunity within a rapidly growing, disruptive start up, please apply More ❯
Posted:
Google JAX
10th Percentile
£62,900
25th Percentile
£64,250
Median
£79,432
75th Percentile
£91,202