diff --git a/src/AbstractOperations/conditional_operations.jl b/src/AbstractOperations/conditional_operations.jl
index 91a110f15f..6ea154e285 100644
--- a/src/AbstractOperations/conditional_operations.jl
+++ b/src/AbstractOperations/conditional_operations.jl
@@ -1,6 +1,6 @@
 using Oceananigans.Fields: OneField
 using Oceananigans.Grids: architecture
-using Oceananigans.Architectures: on_architecture
+import Oceananigans.Architectures: on_architecture
 import Oceananigans.Fields: condition_operand, conditional_length, set!, compute_at!, indices
 
 # For conditional reductions such as mean(u * v, condition = u .> 0))
diff --git a/src/DistributedComputations/distributed_architectures.jl b/src/DistributedComputations/distributed_architectures.jl
index 73ab013355..d8bdc533fb 100644
--- a/src/DistributedComputations/distributed_architectures.jl
+++ b/src/DistributedComputations/distributed_architectures.jl
@@ -204,16 +204,19 @@ Keyword arguments
                   if not for testing or developing. Change at your own risk!
 """
 function Distributed(child_architecture = CPU(); 
-                     communicator = MPI.COMM_WORLD,
+                     communicator = nothing,
                      devices = nothing, 
                      synchronized_communication = false,
-                     partition = Partition(MPI.Comm_size(communicator)))
+                     partition = nothing)
 
     if !(MPI.Initialized())
         @info "MPI has not been initialized, so we are calling MPI.Init()."
         MPI.Init()
     end
 
+    communicator = isnothing(communicator) ? MPI.COMM_WORLD : communicator
+    partition    = isnothing(partition) ? Partition(MPI.Comm_size(communicator)) : partition
+
     ranks = size(partition)
     Rx, Ry, Rz = ranks
     total_ranks = Rx * Ry * Rz
diff --git a/src/DistributedComputations/distributed_on_architecture.jl b/src/DistributedComputations/distributed_on_architecture.jl
index 1564547a59..e02c0172a1 100644
--- a/src/DistributedComputations/distributed_on_architecture.jl
+++ b/src/DistributedComputations/distributed_on_architecture.jl
@@ -1,5 +1,6 @@
 using CUDA: CuArray
 using OffsetArrays
+import Oceananigans.Architectures: on_architecture
 
 # We do not support switching from distributed and serial through `on_architecture`.
 # We only support moving a type from CPU to GPU and the other way around
diff --git a/src/DistributedComputations/partition_assemble.jl b/src/DistributedComputations/partition_assemble.jl
index d381743cb1..a6e730cf47 100644
--- a/src/DistributedComputations/partition_assemble.jl
+++ b/src/DistributedComputations/partition_assemble.jl
@@ -1,4 +1,4 @@
-using Oceananigans.Architectures: on_architecture
+import Oceananigans.Architectures: on_architecture
 
 all_reduce(op, val, arch::Distributed) = 
     MPI.Allreduce(val, op, arch.communicator)
diff --git a/src/ImmersedBoundaries/grid_fitted_bottom.jl b/src/ImmersedBoundaries/grid_fitted_bottom.jl
index 6e9db6bf7e..f456476f61 100644
--- a/src/ImmersedBoundaries/grid_fitted_bottom.jl
+++ b/src/ImmersedBoundaries/grid_fitted_bottom.jl
@@ -4,7 +4,6 @@ using OffsetArrays: OffsetArray
 using Oceananigans.Utils: getnamewrapper
 using Oceananigans.Grids: total_size
 using Oceananigans.Fields: fill_halo_regions!
-using Oceananigans.Architectures: on_architecture
 using Oceananigans.BoundaryConditions: FBC
 using Printf
 
diff --git a/src/ImmersedBoundaries/mask_immersed_field.jl b/src/ImmersedBoundaries/mask_immersed_field.jl
index 4bd2878eec..4f8f0afbad 100644
--- a/src/ImmersedBoundaries/mask_immersed_field.jl
+++ b/src/ImmersedBoundaries/mask_immersed_field.jl
@@ -1,6 +1,5 @@
 using KernelAbstractions: @kernel, @index
 using Statistics
-using Oceananigans.Architectures: architecture
 using Oceananigans.Fields: location, ZReducedField, Field
 
 instantiate(T::Type) = T()
diff --git a/src/ImmersedBoundaries/partial_cell_bottom.jl b/src/ImmersedBoundaries/partial_cell_bottom.jl
index f6ba0368e3..f00771e834 100644
--- a/src/ImmersedBoundaries/partial_cell_bottom.jl
+++ b/src/ImmersedBoundaries/partial_cell_bottom.jl
@@ -1,6 +1,5 @@
 using Oceananigans.Utils: prettysummary
 using Oceananigans.Fields: fill_halo_regions!
-using Oceananigans.Architectures: on_architecture
 using Printf
 
 #####
diff --git a/src/TurbulenceClosures/implicit_explicit_time_discretization.jl b/src/TurbulenceClosures/implicit_explicit_time_discretization.jl
index 33fe206185..3fa9d7a54e 100644
--- a/src/TurbulenceClosures/implicit_explicit_time_discretization.jl
+++ b/src/TurbulenceClosures/implicit_explicit_time_discretization.jl
@@ -1,4 +1,3 @@
-using Oceananigans.Utils: on_architecture
 using Oceananigans.Grids: AbstractGrid
 
 abstract type AbstractTimeDiscretization end
diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl
index a973edc0a0..cfdd407523 100644
--- a/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl
+++ b/src/TurbulenceClosures/turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl
@@ -1,4 +1,4 @@
-using Oceananigans.Architectures: architecture, on_architecture
+using Oceananigans.Architectures: architecture
 using Oceananigans.AbstractOperations: KernelFunctionOperation
 using Oceananigans.BuoyancyModels: ∂z_b
 using Oceananigans.Operators: ℑzᵃᵃᶜ
diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl
index 21f44f3000..0e83f3aa4e 100644
--- a/src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl
+++ b/src/TurbulenceClosures/turbulence_closure_implementations/ri_based_vertical_diffusivity.jl
@@ -1,4 +1,4 @@
-using Oceananigans.Architectures: architecture, on_architecture
+using Oceananigans.Architectures: architecture
 using Oceananigans.BuoyancyModels: ∂z_b
 using Oceananigans.Operators
 using Oceananigans.Grids: inactive_node